2 files changed, 27 insertions, 7 deletions

diff --git a/lib/__set__.js b/lib/__set__.js
index 3fcfbdb..83671a1 100644
--- a/lib/__set__.js
+++ b/lib/__set__.js
@@ -13,18 +13,19 @@ function __set__() {
     arguments.varName = arguments[0];
     arguments.varValue = arguments[1];
     arguments.src = "";
-    arguments.snapshot = {};
+    arguments.revertArgs = [];
 
     if (typeof arguments[0] === "object" && arguments.length === 1) {
         arguments.env = arguments.varName;
         if (!arguments.env || Array.isArray(arguments.env)) {
             throw new TypeError("__set__ expects an object as env");
         }
+        arguments.revertArgs[0] = {};
         for (arguments.varName in arguments.env) {
             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);
+                arguments.revertArgs[0][arguments.varName] = eval(arguments.varName);
             }
         }
     } else if (typeof arguments.varName === "string" && arguments.length === 2) {
@@ -32,16 +33,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);
+        arguments.revertArgs = [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);
+    return function (revertArgs) {
+        __set__.apply(null, revertArgs);
+    }.bind(null, arguments.revertArgs);
 }
 
 module.exports = __set__;
diff --git a/test/testModules/sharedTestCases.js b/test/testModules/sharedTestCases.js
index 20e017d..d482199 100644
--- a/test/testModules/sharedTestCases.js
+++ b/test/testModules/sharedTestCases.js
@@ -267,4 +267,23 @@ describe("rewire " + (typeof testEnv === "undefined"? "(node)": "(" + testEnv +
         }).to.throwException(checkForTypeError);
     });
 
-});
\ No newline at end of file
+    it("should also revert nested changes (with dot notation)", function () {
+        var rewiredModuleA = rewire("./moduleA.js"),
+            revert;
+
+        revert = rewiredModuleA.__set__("myObj.test", true);
+        expect(rewiredModuleA.getMyObj()).to.eql({
+            test: true
+        });
+        revert();
+        // This test also demonstrates a known drawback of the current implementation
+        // If the value doesn't exist at the time it is about to be set, it will be
+        // reverted to undefined instead deleting it from the object
+        // However, this is probably not a real world use-case because why would you
+        // want to mock something when it is not set.
+        expect(rewiredModuleA.getMyObj()).to.eql({
+            test: undefined
+        });
+    });
+
+});