Support an object property function

enzymejs · Nov 6, 2017 · 368899a · 368899a
1 parent b2cc619
commit 368899a
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 13 deletions.
diff --git a/packages/enzyme-test-suite/test/Utils-spec.jsx b/packages/enzyme-test-suite/test/Utils-spec.jsx
@@ -555,15 +555,27 @@ describe('Utils', () => {
           return this.count;
         }
       }
-      const counter = new Counter();
-      const original = counter.incrementAndGet;
-      const spy = spyMethod(counter, 'incrementAndGet');
-      counter.incrementAndGet();
-      counter.incrementAndGet();
-      expect(spy.lastReturnValue).to.equal(3);
-      spy.restore();
-      expect(counter.incrementAndGet).to.equal(original);
-      expect(counter.incrementAndGet()).to.equal(4);
+      const instance = new Counter();
+      const obj = {
+        count: 1,
+        incrementAndGet() {
+          this.count = this.count + 1;
+          return this.count;
+        },
+      };
+
+      // test an instance method and an object property function
+      const targets = [instance, obj];
+      targets.forEach((target) => {
+        const original = target.incrementAndGet;
+        const spy = spyMethod(target, 'incrementAndGet');
+        target.incrementAndGet();
+        target.incrementAndGet();
+        expect(spy.lastReturnValue).to.equal(3);
+        spy.restore();
+        expect(target.incrementAndGet).to.equal(original);
+        expect(target.incrementAndGet()).to.equal(4);
+      });
     });
   });
 });
diff --git a/packages/enzyme/src/Utils.js b/packages/enzyme/src/Utils.js
@@ -269,20 +269,28 @@ export function cloneElement(adapter, el, props) {
 
 export function spyMethod(instance, methodName) {
   let lastReturnValue;
+  const originalMethod = instance[methodName];
+  const hasOwn = Object.hasOwnProperty.call(instance, methodName);
   Object.defineProperty(instance, methodName, {
     configurable: true,
     enumerable: false,
     value(...args) {
-      const result = this.constructor.prototype[methodName].apply(this, args);
+      const result = originalMethod.apply(this, args);
       lastReturnValue = result;
       return result;
     },
   });
   return {
     restore() {
-      /* eslint-disable no-param-reassign */
-      delete instance[methodName];
-      /* eslint-enable no-param-reassign */
+      if (hasOwn) {
+        /* eslint-disable no-param-reassign */
+        instance[methodName] = originalMethod;
+        /* eslint-enable no-param-reassign */
+      } else {
+        /* eslint-disable no-param-reassign */
+        delete instance[methodName];
+        /* eslint-enable no-param-reassign */
+      }
     },
     get lastReturnValue() {
       return lastReturnValue;