2 files changed, 79 insertions, 3 deletions

diff --git a/lib/__with__.js b/lib/__with__.js
index ef29c8c..387714d 100644
--- a/lib/__with__.js
+++ b/lib/__with__.js
@@ -15,7 +15,9 @@ function __with__() {
     var args = arguments;
 
     return function (callback) {
-        var undo;
+        var undo,
+            returned,
+            isPromise;
 
         if (typeof callback !== "function") {
             throw new TypeError("__with__ expects a callback function");
@@ -24,9 +26,16 @@ function __with__() {
         undo = module.exports.__set__.apply(null, args);
 
         try {
-            callback();
+            returned = callback();
+            isPromise = returned && typeof returned.then === "function";
+            if (isPromise) {
+                returned.then(undo, undo);
+                return returned;
+            }
         } finally {
-            undo();
+            if (!isPromise) {
+                undo();
+            }
         }
     };
 }
diff --git a/test/__with__.test.js b/test/__with__.test.js
index 532ef91..3cd4b57 100644
--- a/test/__with__.test.js
+++ b/test/__with__.test.js
@@ -111,4 +111,71 @@ describe("__with__", function() {
         expect(callWithFunction({})).to.throwError(expectTypeError);
         expect(callWithFunction(function(){})).to.not.throwError(expectTypeError);
     });
+
+    describe("using promises", function () {
+        var promiseFake;
+
+        beforeEach(function () {
+            promiseFake = {
+                then: function (onResolve, onReject) {
+                    promiseFake.onResolve = onResolve;
+                    promiseFake.onReject = onReject;
+                }
+            };
+        });
+
+        it("should pass the returned promise through", function () {
+            var fn = moduleFake.__with__({});
+
+            expect(fn(function () {
+                return promiseFake;
+            })).to.equal(promiseFake);
+        });
+
+        it("should not undo any changes until the promise has been resolved", function () {
+            expect(moduleFake.getValue()).to.be(0);
+            expect(moduleFake.getReference()).to.eql({});
+
+            moduleFake.__with__({
+                myValue: 2,
+                myReference: newObj
+            })(function () {
+                return promiseFake;
+            });
+
+            // the change should still be present at this point
+            expect(moduleFake.getValue()).to.be(2);
+            expect(moduleFake.getReference()).to.be(newObj);
+
+            promiseFake.onResolve();
+
+            // now everything should be back to normal
+            expect(moduleFake.getValue()).to.be(0);
+            expect(moduleFake.getReference()).to.eql({});
+        });
+
+        it("should also undo any changes if the promise has been rejected", function () {
+            expect(moduleFake.getValue()).to.be(0);
+            expect(moduleFake.getReference()).to.eql({});
+
+            moduleFake.__with__({
+                myValue: 2,
+                myReference: newObj
+            })(function () {
+                return promiseFake;
+            });
+
+            // the change should still be present at this point
+            expect(moduleFake.getValue()).to.be(2);
+            expect(moduleFake.getReference()).to.be(newObj);
+
+            promiseFake.onReject();
+
+            // now everything should be back to normal
+            expect(moduleFake.getValue()).to.be(0);
+            expect(moduleFake.getReference()).to.eql({});
+        });
+
+    });
+
 });