3 files changed, 102 insertions, 5 deletions

diff --git a/lib/__set__.js b/lib/__set__.js
index 0fa88ac..ad0e5b0 100644
--- a/lib/__set__.js
+++ b/lib/__set__.js
@@ -46,4 +46,21 @@ function __set__() {
     };
 }
 
-module.exports = __set__;
+function __with__() {
+  var args = arguments;
+  return function(callback) {
+    if (typeof callback !== "function") {
+      throw new TypeError("__with__ expects a callback function")
+    }
+
+    var undo = __set__.apply(null, args)
+    try {
+      callback();
+    }
+    finally {
+      undo();
+    }
+  }
+}
+
+module.exports = {"__set__": __set__, "__with__": __with__}
diff --git a/lib/rewire.js b/lib/rewire.js
index 3daacef..61160fd 100644
--- a/lib/rewire.js
+++ b/lib/rewire.js
@@ -1,13 +1,16 @@
 var Module = require("module"),
     fs = require("fs"),
     __get__ = require("./__get__.js"),
-    __set__ = require("./__set__.js"),
+    setModule = require ("./__set__.js"),
+    __set__ = setModule["__set__"],
+    __with__ = setModule["__with__"],
     getImportGlobalsSrc = require("./getImportGlobalsSrc.js"),
     detectStrictMode = require("./detectStrictMode.js"),
     moduleEnv = require("./moduleEnv.js");
 
 var __get__Src = __get__.toString(),
-    __set__Src = __set__.toString();
+    __set__Src = __set__.toString(),
+    __with_Src = __with__.toString();
 
 /**
  * Does actual rewiring the module. For further documentation @see index.js
@@ -44,6 +47,7 @@ function internalRewire(parentModulePath, targetPath) {
     appendix = "\n";
     appendix += "module.exports.__set__ = " + __set__Src + "; ";
     appendix += "module.exports.__get__ = " + __get__Src + "; ";
+    appendix += "module.exports.__with__ = " + __with_Src + "; ";
 
     // Check if the module uses the strict mode.
     // If so we must ensure that "use strict"; stays at the beginning of the module.
@@ -58,4 +62,4 @@ function internalRewire(parentModulePath, targetPath) {
     return targetModule.exports;
 }
 
-module.exports = internalRewire;
\ No newline at end of file
+module.exports = internalRewire;
diff --git a/test/__set__.test.js b/test/__set__.test.js
index ff1a3d2..fe6f529 100644
--- a/test/__set__.test.js
+++ b/test/__set__.test.js
@@ -1,5 +1,7 @@
 var expect = require("expect.js"),
-    __set__ = require("../lib/__set__.js"),
+    setModule = require("../lib/__set__.js")
+    __set__ = setModule["__set__"],
+    __with__ = setModule["__with__"],
     vm = require("vm"),
 
     expectReferenceError = expectError(ReferenceError),
@@ -125,3 +127,77 @@ describe("__set__", function () {
         }).to.throwException(expectTypeError);
     });
 });
+
+describe("__with__", function() {
+    var moduleFake;
+
+    beforeEach(function () {
+        moduleFake = {
+            myValue: 0,    // copy by value
+            myReference: {}       // copy by reference
+        };
+
+        //__with__ requires __set__ to be in scope
+        vm.runInNewContext(
+            "__set__ = " + __set__.toString() + "; " +
+            "__with__ = " + __with__.toString() + "; " +
+            "getValue = function () { return myValue; }; " +
+            "getReference = function () { return myReference; }; ",
+            moduleFake
+        );
+    });
+
+    it("should return a function that can be invoked with a callback which guarantees __sets__ undo function is called for you at the end", function () {
+        var newObj = { hello: "hello" };
+
+        expect(moduleFake.getValue()).to.be(0);
+        expect(moduleFake.getReference()).to.eql({});
+
+        moduleFake.__with__({
+            myValue: 2,
+            myReference: newObj
+        })(function() {
+          //changes will be visible from within this callback function
+          expect(moduleFake.getValue()).to.be(2);
+          expect(moduleFake.getReference()).to.be(newObj);
+        })
+
+        //undo will automatically get called for you after returning from your callback function
+        expect(moduleFake.getValue()).to.be(0);
+        expect(moduleFake.getReference()).to.eql({});
+    });
+
+    it("should still revert values if the callback throws an exception", function(){
+        var newObj = { hello: "hello" };
+        function withError(){
+          moduleFake.__with__({
+              myValue: 2,
+              myReference: newObj
+          })(function() {
+            throw new Error("something went wrong...");
+          })
+        }
+        expect(withError).to.throwError();
+        expect(moduleFake.getValue()).to.be(0);
+        expect(moduleFake.getReference()).to.eql({});
+    });
+
+    it("should throw an error if something other than a function is passed as the callback", function() {
+        var newObj = { hello: "hello" },
+            withFunction = moduleFake.__with__({
+                myValue: 2,
+                myReference: newObj
+            })
+            callWithFunction = function(){
+              var args = arguments;
+              return function() {
+                withFunction.apply(null, args);
+              };
+            };
+
+        expect(callWithFunction(1)).to.throwError();
+        expect(callWithFunction("a string")).to.throwError();
+        expect(callWithFunction({})).to.throwError();
+        expect(callWithFunction(function(){})).to.not.throwError();
+    });
+});