2 files changed, 34 insertions, 4 deletions

diff --git a/lib/__set__.js b/lib/__set__.js
index 66e8c0a..0fa88ac 100644
--- a/lib/__set__.js
+++ b/lib/__set__.js
@@ -11,10 +11,12 @@
  * @throws {ReferenceError} When the variable is unknown
  * @return {*}
  */
+
 function __set__() {
     arguments.varName = arguments[0];
     arguments.varValue = arguments[1];
     arguments.src = "";
+    var snapshot = {};
 
     if (typeof arguments[0] === "object" && arguments.length === 1) {
         arguments.env = arguments.varName;
@@ -25,6 +27,7 @@ function __set__() {
             if (arguments.env.hasOwnProperty(arguments.varName)) {
                 arguments.varValue = arguments.env[arguments.varName];
                 arguments.src += arguments.varName + " = arguments.env." + arguments.varName + "; ";
+                snapshot[arguments.varName] = eval(arguments.varName);
             }
         }
     } else if (typeof arguments.varName === "string" && arguments.length === 2) {
@@ -32,11 +35,15 @@ function __set__() {
             throw new TypeError("__set__ expects a non-empty string as a variable name");
         }
         arguments.src = arguments.varName + " = arguments.varValue;";
+        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() {
+      __set__(snapshot);
+    };
 }
 
-module.exports = __set__;
\ No newline at end of file
+module.exports = __set__;
diff --git a/test/__set__.test.js b/test/__set__.test.js
index 7b9c2cb..ff1a3d2 100644
--- a/test/__set__.test.js
+++ b/test/__set__.test.js
@@ -69,8 +69,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(typeof undo).to.be("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({});
+
+        var 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 () {
@@ -101,4 +124,4 @@ describe("__set__", function () {
             moduleFake.__set__("someVar");  // misfitting number of params
         }).to.throwException(expectTypeError);
     });
-});
\ No newline at end of file
+});