Implement spyMethod and use it to inspect the result of shouldComponentUpdate

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);
+    });
+  });
 });

packages/enzyme/src/ShallowWrapper.js

@@ -18,6 +18,7 @@ import {
   sym,
   privateSet,
   cloneElement,
+  spyMethod,
 } from './Utils';
 import { debugNodes } from './Debug';
 import {
@@ -310,78 +311,61 @@ class ShallowWrapper {
         const { state } = instance;
         const prevProps = instance.props || this[UNRENDERED].props;
         const prevContext = instance.context || this[OPTIONS].context;
-        const nextProps = { ...prevProps, ...props };
         const nextContext = context || prevContext;
         if (context) {
           this[OPTIONS] = { ...this[OPTIONS], context: nextContext };
         }
         this[RENDERER].batchedUpdates(() => {
+          // When shouldComponentUpdate returns false we shouldn't call componentDidUpdate.
+          // so we spy shouldComponentUpdate to get the result.
           let shouldRender = true;
-          // dirty hack:
-          // make sure that componentWillReceiveProps is called before shouldComponentUpdate
-          let originalComponentWillReceiveProps;
+          let spy;
           if (
             !this[OPTIONS].disableLifecycleMethods &&
             instance &&
-            typeof instance.componentWillReceiveProps === 'function'
+            typeof instance.shouldComponentUpdate === 'function'
           ) {
-            instance.componentWillReceiveProps(nextProps, nextContext);
-            originalComponentWillReceiveProps = instance.componentWillReceiveProps;
-            instance.componentWillReceiveProps = () => {};
+            spy = spyMethod(instance, 'shouldComponentUpdate');
+          }
+          if (props) this[UNRENDERED] = cloneElement(adapter, this[UNRENDERED], props);
+          this[RENDERER].render(this[UNRENDERED], nextContext);
+          if (spy) {
+            shouldRender = spy.getLastReturnValue();
+            spy.restore();
           }
-          // dirty hack: avoid calling shouldComponentUpdate twice
-          let originalShouldComponentUpdate;
           if (
+            shouldRender &&
             !this[OPTIONS].disableLifecycleMethods &&
-            instance &&
-            typeof instance.shouldComponentUpdate === 'function'
+            instance
           ) {
-            shouldRender = instance.shouldComponentUpdate(nextProps, state, nextContext);
-            originalShouldComponentUpdate = instance.shouldComponentUpdate;
-          }
-          if (shouldRender) {
-            if (props) this[UNRENDERED] = cloneElement(adapter, this[UNRENDERED], props);
-            if (originalShouldComponentUpdate) {
-              instance.shouldComponentUpdate = () => true;
-            }
-
-            this[RENDERER].render(this[UNRENDERED], nextContext);
-
-            if (originalShouldComponentUpdate) {
-              instance.shouldComponentUpdate = originalShouldComponentUpdate;
-            }
             const lifecycles = getAdapterLifecycles(adapter);
-            if (
-              !this[OPTIONS].disableLifecycleMethods &&
-              instance
-            ) {
+
+            if (lifecycles.getSnapshotBeforeUpdate) {
+              let snapshot;
+              if (typeof instance.getSnapshotBeforeUpdate === 'function') {
+                snapshot = instance.getSnapshotBeforeUpdate(prevProps, state);
+              }
               if (
-                lifecycles.getSnapshotBeforeUpdate
-                && typeof instance.getSnapshotBeforeUpdate === 'function'
+                lifecycles.componentDidUpdate &&
+                typeof instance.componentDidUpdate === 'function'
               ) {
-                const snapshot = instance.getSnapshotBeforeUpdate(prevProps, state);
-                if (typeof instance.componentDidUpdate === 'function') {
-                  instance.componentDidUpdate(prevProps, state, snapshot);
-                }
-              } else if (typeof instance.componentDidUpdate === 'function') {
-                if (
-                  lifecycles.componentDidUpdate &&
-                  lifecycles.componentDidUpdate.prevContext
-                ) {
-                  instance.componentDidUpdate(prevProps, state, prevContext);
-                } else {
-                  instance.componentDidUpdate(prevProps, state);
-                }
+                instance.componentDidUpdate(prevProps, state, snapshot);
+              }
+            } else if (
+              lifecycles.componentDidUpdate &&
+              typeof instance.componentDidUpdate === 'function'
+            ) {
+              if (lifecycles.componentDidUpdate.prevContext) {
+                instance.componentDidUpdate(prevProps, state, prevContext);
+              } else {
+                instance.componentDidUpdate(prevProps, state);
               }
             }
-            this.update();
           // If it doesn't need to rerender, update only its props.
           } else if (props) {
             instance.props = props;
           }
-          if (originalComponentWillReceiveProps) {
-            instance.componentWillReceiveProps = originalComponentWillReceiveProps;
-          }
+          this.update();
         });
       });
     });
@@ -438,30 +422,26 @@ class ShallowWrapper {
         const prevProps = instance.props;
         const prevState = instance.state;
         const prevContext = instance.context;
-        let shouldRender = true;
-        // This is a dirty hack but it requires to know the result of shouldComponentUpdate.
         // When shouldComponentUpdate returns false we shouldn't call componentDidUpdate.
-        // shouldComponentUpdate is called in `instance.setState`
-        // so we replace shouldComponentUpdate to know the result and restore it later.
-        let originalShouldComponentUpdate;
+        // so we spy shouldComponentUpdate to get the result.
+        let spy;
+        let shouldRender = true;
         if (
           !this[OPTIONS].disableLifecycleMethods &&
           lifecycles.componentDidUpdate &&
           lifecycles.componentDidUpdate.onSetState &&
           instance &&
           typeof instance.shouldComponentUpdate === 'function'
         ) {
-          originalShouldComponentUpdate = instance.shouldComponentUpdate;
-          instance.shouldComponentUpdate = (...args) => {
-            shouldRender = originalShouldComponentUpdate.apply(instance, args);
-            instance.shouldComponentUpdate = originalShouldComponentUpdate;
-            return shouldRender;
-          };
+          spy = spyMethod(instance, 'shouldComponentUpdate');
         }
         // We don't pass the setState callback here
         // to guarantee to call the callback after finishing the render
         instance.setState(state);
-
+        if (spy) {
+          shouldRender = spy.getLastReturnValue();
+          spy.restore();
+        }
         if (
           shouldRender &&
           !this[OPTIONS].disableLifecycleMethods &&

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;
+    },
+  };
+}