Merge branch 'bobpace-undo'

6 files changed, 278 insertions, 15 deletions

diff --git a/.istanbul.yml b/.istanbul.yml
index 480671f..21bf908 100644
--- a/.istanbul.yml
+++ b/.istanbul.yml
@@ -1,3 +1,7 @@
 instrumentation:
-    # __get__ and __set__ will be stringified and evaled again. Thus it's difficult to include them into the test coverage
-    excludes: ['lib/__get__.js', 'lib/__set__.js']
\ No newline at end of file
+    # These functions will be stringified and evaled again. Thus it's difficult to include them into the test coverage
+    excludes: [
+        'lib/__get__.js',
+        'lib/__set__.js',
+        'lib/__with__.js'
+        ]
\ No newline at end of file
diff --git a/lib/__set__.js b/lib/__set__.js
index 66e8c0a..3fcfbdb 100644
--- a/lib/__set__.js
+++ b/lib/__set__.js
@@ -5,16 +5,15 @@
  * All variables within this function are namespaced in the arguments array because every
  * var declaration could possibly clash with a variable in the module scope.
  *
- * @param {!String|!Object} varName name of the variable to set
+ * @param {String|Object} varName name of the variable to set
  * @param {String} varValue new value
- * @throws {TypeError}
- * @throws {ReferenceError} When the variable is unknown
- * @return {*}
+ * @return {Function}
  */
 function __set__() {
     arguments.varName = arguments[0];
     arguments.varValue = arguments[1];
     arguments.src = "";
+    arguments.snapshot = {};
 
     if (typeof arguments[0] === "object" && arguments.length === 1) {
         arguments.env = arguments.varName;
@@ -25,6 +24,7 @@ function __set__() {
             if (arguments.env.hasOwnProperty(arguments.varName)) {
                 arguments.varValue = arguments.env[arguments.varName];
                 arguments.src += arguments.varName + " = arguments.env." + arguments.varName + "; ";
+                arguments.snapshot[arguments.varName] = eval(arguments.varName);
             }
         }
     } else if (typeof arguments.varName === "string" && arguments.length === 2) {
@@ -32,11 +32,16 @@ function __set__() {
             throw new TypeError("__set__ expects a non-empty string as a variable name");
         }
         arguments.src = arguments.varName + " = arguments.varValue;";
+        arguments.snapshot[arguments.varName] = eval(arguments.varName);
     } else {
         throw new TypeError("__set__ expects an environment object or a non-empty string as a variable name");
     }
 
     eval(arguments.src);
+
+    return function (snapshot) {
+        module.exports.__set__(snapshot);
+    }.bind(null, arguments.snapshot);
 }
 
-module.exports = __set__;
\ No newline at end of file
+module.exports = __set__;
diff --git a/lib/__with__.js b/lib/__with__.js
new file mode 100644
index 0000000..387714d
--- /dev/null
+++ b/lib/__with__.js
@@ -0,0 +1,43 @@
+"use strict";
+
+/**
+ * This function will be stringified and then injected into every rewired module.
+ *
+ * Calling myModule.__with__("myPrivateVar", newValue) returns a function where
+ * you can place your tests. As long as the returned function is executed variables
+ * will be set to the given value, after that all changed variables are reset back to normal.
+ *
+ * @param {String|Object} varName name of the variable to set
+ * @param {String} varValue new value
+ * @return {Function}
+ */
+function __with__() {
+    var args = arguments;
+
+    return function (callback) {
+        var undo,
+            returned,
+            isPromise;
+
+        if (typeof callback !== "function") {
+            throw new TypeError("__with__ expects a callback function");
+        }
+
+        undo = module.exports.__set__.apply(null, args);
+
+        try {
+            returned = callback();
+            isPromise = returned && typeof returned.then === "function";
+            if (isPromise) {
+                returned.then(undo, undo);
+                return returned;
+            }
+        } finally {
+            if (!isPromise) {
+                undo();
+            }
+        }
+    };
+}
+
+module.exports = __with__;
\ No newline at end of file
diff --git a/lib/rewire.js b/lib/rewire.js
index 3daacef..086cc27 100644
--- a/lib/rewire.js
+++ b/lib/rewire.js
@@ -1,13 +1,15 @@
 var Module = require("module"),
     fs = require("fs"),
     __get__ = require("./__get__.js"),
-    __set__ = require("./__set__.js"),
+    __set__ = require ("./__set__.js"),
+    __with__ = require("./__with__.js"),
     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 +46,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 +61,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 7b9c2cb..ef69a07 100644
--- a/test/__set__.test.js
+++ b/test/__set__.test.js
@@ -2,7 +2,6 @@ var expect = require("expect.js"),
     __set__ = require("../lib/__set__.js"),
     vm = require("vm"),
 
-    expectReferenceError = expectError(ReferenceError),
     expectTypeError = expectError(TypeError);
 
 function expectError(ErrConstructor) {
@@ -12,16 +11,21 @@ function expectError(ErrConstructor) {
 }
 
 describe("__set__", function () {
-    var moduleFake;
+    var moduleFake,
+        undo;
 
     beforeEach(function () {
         moduleFake = {
+            module: {
+                exports: {}
+            },
             myValue: 0,    // copy by value
             myReference: {}       // copy by reference
         };
 
         vm.runInNewContext(
-            "__set__ = " + __set__.toString() + "; " +
+            //__set__ requires __set__ to be present on module.exports
+            "__set__ = module.exports.__set__ = " + __set__.toString() + "; " +
             "getValue = function () { return myValue; }; " +
             "getReference = function () { return myReference; }; ",
             moduleFake
@@ -69,8 +73,31 @@ describe("__set__", function () {
         expect(moduleFake.getValue()).to.be(2);
         expect(moduleFake.getReference()).to.be(newObj);
     });
-    it("should return undefined", function () {
-        expect(moduleFake.__set__("myValue", 4)).to.be(undefined);
+    it("should return a function that when invoked reverts to the values before set was called", function () {
+        undo = moduleFake.__set__("myValue", 4);
+        expect(undo).to.be.a("function");
+        expect(moduleFake.getValue()).to.be(4);
+        undo();
+        expect(moduleFake.getValue()).to.be(0);
+    });
+    it("should be able to revert when calling with an env-obj", function () {
+        var newObj = { hello: "hello" };
+
+        expect(moduleFake.getValue()).to.be(0);
+        expect(moduleFake.getReference()).to.eql({});
+
+        undo = moduleFake.__set__({
+            myValue: 2,
+            myReference: newObj
+        });
+
+        expect(moduleFake.getValue()).to.be(2);
+        expect(moduleFake.getReference()).to.be(newObj);
+
+        undo();
+
+        expect(moduleFake.getValue()).to.be(0);
+        expect(moduleFake.getReference()).to.eql({});
     });
     it("should throw a TypeError when passing misfitting params", function () {
         expect(function () {
diff --git a/test/__with__.test.js b/test/__with__.test.js
new file mode 100644
index 0000000..3cd4b57
--- /dev/null
+++ b/test/__with__.test.js
@@ -0,0 +1,181 @@
+var expect = require("expect.js"),
+    __with__ = require("../lib/__with__.js"),
+    __set__ = require("../lib/__set__.js"),
+    vm = require("vm"),
+
+    expectTypeError = expectError(TypeError);
+
+function expectError(ErrConstructor) {
+    return function expectReferenceError(err) {
+        expect(err.constructor.name).to.be(ErrConstructor.name);
+    };
+}
+
+describe("__with__", function() {
+    var moduleFake,
+        newObj;
+
+    beforeEach(function () {
+        moduleFake = {
+            module: {
+                exports: {}
+            },
+            myValue: 0,    // copy by value
+            myReference: {}       // copy by reference
+        };
+
+        newObj = { hello: "hello" };
+
+        vm.runInNewContext(
+            //__with__ requires __set__ to be present on module.exports
+            "module.exports.__set__ = " + __set__.toString() + "; " +
+            "__with__ = " + __with__.toString() + "; " +
+            "getValue = function () { return myValue; }; " +
+            "getReference = function () { return myReference; }; ",
+            moduleFake
+        );
+    });
+
+    it("should return a function", function () {
+        expect(moduleFake.__with__({
+            myValue: 2,
+            myReference: newObj
+        })).to.be.a("function");
+    });
+
+    it("should return a function that can be invoked with a callback which guarantees __set__'s undo function is called for you at the end", function () {
+        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 also accept a variable name and a variable value (just like __set__)", function () {
+        expect(moduleFake.getValue()).to.be(0);
+
+        moduleFake.__with__("myValue", 2)(function () {
+            expect(moduleFake.getValue()).to.be(2);
+        });
+
+        expect(moduleFake.getValue()).to.be(0);
+
+        expect(moduleFake.getReference()).to.eql({});
+
+        moduleFake.__with__("myReference", newObj)(function () {
+            expect(moduleFake.getReference()).to.be(newObj);
+        });
+
+        expect(moduleFake.getReference()).to.eql({});
+    });
+
+    it("should still revert values if the callback throws an exception", function(){
+        expect(function withError() {
+            moduleFake.__with__({
+                myValue: 2,
+                myReference: newObj
+            })(function () {
+                throw new Error("something went wrong...");
+            });
+        }).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 withFunction = moduleFake.__with__({
+                myValue: 2,
+                myReference: newObj
+            });
+
+        function callWithFunction() {
+            var args = arguments;
+
+            return function () {
+                withFunction.apply(null, args);
+            };
+        }
+
+        expect(callWithFunction(1)).to.throwError(expectTypeError);
+        expect(callWithFunction("a string")).to.throwError(expectTypeError);
+        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({});
+        });
+
+    });
+
+});