[New] add spyMethod

enzymejs · Sep 28, 2017 · 9d6e304 · 9d6e304
1 parent eb55307
commit 9d6e304
Show file tree

Hide file tree

Showing 2 changed files with 90 additions and 0 deletions.
diff --git a/packages/enzyme-test-suite/test/Utils-spec.jsx b/packages/enzyme-test-suite/test/Utils-spec.jsx
@@ -5,6 +5,7 @@ import {
   nodeEqual,
   nodeMatches,
   displayNameOfNode,
+  spyMethod,
 } from 'enzyme/build/Utils';
 import {
   flatten,
@@ -551,4 +552,53 @@ describe('Utils', () => {
       expect(flat).to.deep.equal([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
     });
   });
+
+  describe('spyMethod', () => {
+    it('should be able to spy last return value and restore it', () => {
+      class Counter {
+        constructor() {
+          this.count = 1;
+        }
+        incrementAndGet() {
+          this.count = this.count + 1;
+          return this.count;
+        }
+      }
+      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.getLastReturnValue()).to.equal(3);
+        spy.restore();
+        expect(target.incrementAndGet).to.equal(original);
+        expect(target.incrementAndGet()).to.equal(4);
+      });
+    });
+
+    it('should be able to restore the property descriptor', () => {
+      const obj = {};
+      const descriptor = {
+        configurable: true,
+        enumerable: true,
+        writable: true,
+        value: () => {},
+      };
+      Object.defineProperty(obj, 'method', descriptor);
+      const spy = spyMethod(obj, 'method');
+      spy.restore();
+      expect(Object.getOwnPropertyDescriptor(obj, 'method')).to.deep.equal(descriptor);
+    });
+  });
 });
diff --git a/packages/enzyme/src/Utils.js b/packages/enzyme/src/Utils.js
@@ -3,6 +3,7 @@ import isEqual from 'lodash/isEqual';
 import is from 'object-is';
 import entries from 'object.entries';
 import functionName from 'function.prototype.name';
+import has from 'has';
 import configuration from './configuration';
 import validateAdapter from './validateAdapter';
 import { childrenOfNode } from './RSTTraversal';
@@ -242,3 +243,42 @@ export function cloneElement(adapter, el, props) {
     { ...el.props, ...props },
   );
 }
+
+export function spyMethod(instance, methodName) {
+  let lastReturnValue;
+  const originalMethod = instance[methodName];
+  const hasOwn = has(instance, methodName);
+  let descriptor;
+  if (hasOwn) {
+    descriptor = Object.getOwnPropertyDescriptor(instance, methodName);
+  }
+  Object.defineProperty(instance, methodName, {
+    configurable: true,
+    enumerable: !descriptor || !!descriptor.enumerable,
+    value(...args) {
+      const result = originalMethod.apply(this, args);
+      lastReturnValue = result;
+      return result;
+    },
+  });
+  return {
+    restore() {
+      if (hasOwn) {
+        if (descriptor) {
+          Object.defineProperty(instance, methodName, descriptor);
+        } else {
+          /* 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 */
+      }
+    },
+    getLastReturnValue() {
+      return lastReturnValue;
+    },
+  };
+}