Refactor; introduce sinon.wrapMethod to handle wrapping of methods on…

… object with a restore() option

src/sinon.js

@@ -1,3 +1,30 @@
 var sinon = (function () {
-  return {};
+  function wrapMethod (object, property, method) {
+    if (!object) {
+      throw new TypeError("Should wrap property of object");
+    }
+
+    if (typeof method != "function") {
+      throw new TypeError("Method wrapper should be function");
+    }
+
+    var wrappedMethod = object[property];
+    var type = typeof wrappedMethod;
+
+    if (!!wrappedMethod && type != "function") {
+      throw new TypeError("Attempted to wrap " + type + " as function");
+    }
+
+    object[property] = method;
+
+    method.restore = function () {
+      object[property] = wrappedMethod;
+    };
+
+    return method;
+  }
+
+  return {
+    wrapMethod: wrapMethod
+  };
 }());

src/spy.js

@@ -1,4 +1,9 @@
 (function () {
+  function stub (object, property) {
+    var method = object[property];
+    return sinon.wrapMethod(object, property, spy.create(method));
+  }
+
   function spy (func) {
     return spy.create(func);
   }

src/stub.js

@@ -1,31 +1,18 @@
 (function () {
   function stub (object, property, func) {
-    if (!object) {
-      throw new TypeError("Should stub property of object");
-    }
-
     if (!!func && typeof func != "function") {
       throw new TypeError("Custom stub should be function");
     }
 
-    var method = object[property];
-    var type = typeof method;
-
-    if (!!method && type != "function") {
-      throw new TypeError("Attempted to stub " + type + " as function");
-    }
+    var wrapper;
 
-    if (typeof func == "function") {
-      object[property] = !!sinon.spy ? sinon.spy(func) : func;
+    if (func) {
+      wrapper = !!sinon.spy ? sinon.spy(func) : func;
     } else {
-      object[property] = stub.create();
+      wrapper = stub.create();
     }
 
-    object[property].restore = function () {
-      object[property] = method;
-    };
-
-    return object[property];
+    return sinon.wrapMethod(object, property, wrapper);
   }
 
   Object.extend(stub, (function () {