[enzyme-adapter-react-16] [fix] properly handle memo’s areEquals arg

Fixes #2128

packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js

@@ -36,7 +36,7 @@ import {
   Suspense,
 } from 'react-is';
 import { EnzymeAdapter } from 'enzyme';
-import { typeOfNode } from 'enzyme/build/Utils';
+import { typeOfNode, shallowEqual } from 'enzyme/build/Utils';
 import {
   displayNameOfNode,
   elementToTree as utilElementToTree,
@@ -358,6 +358,13 @@ function makeFakeElement(type) {
   return { $$typeof: Element, type };
 }
 
+function isStateful(Component) {
+  return Component.prototype && (
+    Component.prototype.isReactComponent
+    || Array.isArray(Component.__reactAutoBindPairs) // fallback for createClass components
+  );
+}
+
 class ReactSixteenAdapter extends EnzymeAdapter {
   constructor() {
     super();
@@ -497,6 +504,45 @@ class ReactSixteenAdapter extends EnzymeAdapter {
 
     let lastComponent = null;
     let wrappedComponent = null;
+    const sentinel = {};
+
+    // wrap memo components with a PureComponent, or a class component with sCU
+    const wrapPureComponent = (Component, compare) => {
+      if (!is166) {
+        throw new RangeError('this function should not be called in React < 16.6. Please report this!');
+      }
+      if (lastComponent !== Component) {
+        if (isStateful(Component)) {
+          wrappedComponent = class extends Component {}; // eslint-disable-line react/prefer-stateless-function
+          if (compare) {
+            wrappedComponent.prototype.shouldComponentUpdate = nextProps => !compare(this.props, nextProps);
+          } else {
+            wrappedComponent.prototype.isPureReactComponent = true;
+          }
+        } else {
+          let memoized = sentinel;
+          let prevProps;
+          wrappedComponent = function (props, ...args) {
+            const shouldUpdate = memoized === sentinel || (compare
+              ? !compare(prevProps, props)
+              : !shallowEqual(prevProps, props)
+            );
+            if (shouldUpdate) {
+              memoized = Component({ ...Component.defaultProps, ...props }, ...args);
+              prevProps = props;
+            }
+            return memoized;
+          };
+        }
+        Object.assign(
+          wrappedComponent,
+          Component,
+          { displayName: adapter.displayNameOfNode({ type: Component }) },
+        );
+        lastComponent = Component;
+      }
+      return wrappedComponent;
+    };
 
     // Wrap functional components on versions prior to 16.5,
     // to avoid inadvertently pass a `this` instance to it.
@@ -507,6 +553,7 @@ class ReactSixteenAdapter extends EnzymeAdapter {
             // eslint-disable-next-line new-cap
             (props, ...args) => Component({ ...Component.defaultProps, ...props }, ...args),
             Component,
+            { displayName: adapter.displayNameOfNode({ type: Component }) },
           );
           lastComponent = Component;
         }
@@ -567,22 +614,19 @@ class ReactSixteenAdapter extends EnzymeAdapter {
             renderedEl = React.createElement(FakeSuspenseWrapper, null, children);
           }
           const { type: Component } = renderedEl;
-          const isStateful = Component.prototype && (
-            Component.prototype.isReactComponent
-            || Array.isArray(Component.__reactAutoBindPairs) // fallback for createClass components
-          );
 
           const context = getMaskedContext(Component.contextTypes, unmaskedContext);
 
-          if (!isStateful && isMemo(el.type)) {
-            const InnerComp = el.type.type;
+          if (isMemo(el.type)) {
+            const { type: InnerComp, compare } = el.type;
+
             return withSetStateAllowed(() => renderer.render(
-              { ...el, type: wrapFunctionalComponent(InnerComp) },
+              { ...el, type: wrapPureComponent(InnerComp, compare) },
               context,
             ));
           }
 
-          if (!isStateful && typeof Component === 'function') {
+          if (!isStateful(Component) && typeof Component === 'function') {
             return withSetStateAllowed(() => renderer.render(
               { ...renderedEl, type: wrapFunctionalComponent(Component) },
               context,

packages/enzyme-test-suite/test/shared/methods/text.jsx

@@ -11,6 +11,7 @@ import { is } from '../../_helpers/version';
 
 import {
   Fragment,
+  memo,
 } from '../../_helpers/react-compat';
 
 export default function describeText({
@@ -256,5 +257,101 @@ export default function describeText({
         expect(text).to.equal(isShallow ? '<Title />Bar' : 'FooBar');
       });
     });
+
+    describeIf(is('> 16.6'), 'React.memo', () => {
+      class Dumb extends React.Component {
+        render() {
+          const { number } = this.props;
+          return <span>It’s number {String(number)}</span>;
+        }
+      }
+
+      function DumbSFC({ number }) {
+        return <span>It’s number {String(number)}</span>;
+      }
+
+      function areEqual(props, nextProps) {
+        return nextProps.number > 10 && nextProps.number < 20;
+      }
+
+      const DumbMemo = memo && memo(Dumb, areEqual);
+      const DumbMemoNoCompare = memo && memo(Dumb);
+      const DumbSFCMemo = memo && memo(DumbSFC, areEqual);
+      const DumbSFCMemoNoCompare = memo && memo(DumbSFC);
+
+      it('<Dumb /> - should always re-render', () => {
+        const tree = Wrap(<Dumb number={5} />);
+
+        expect(tree.text()).to.equal('It’s number 5');
+
+        tree.setProps({ number: 15 });
+
+        expect(tree.text()).to.equal('It’s number 15');
+      });
+
+      it('<DumbMemo /> - should only re-render when number is between 10 and 20', () => {
+        const tree = Wrap(<DumbMemo number={5} />);
+
+        expect(tree.text()).to.equal('It’s number 5');
+
+        tree.setProps({ number: 15 });
+
+        expect(tree.text()).to.equal('It’s number 5');
+      });
+
+      it('<DumbSFC /> - should always re-render', () => {
+        const tree = Wrap(<DumbSFC number={5} />);
+
+        expect(tree.text()).to.equal('It’s number 5');
+
+        tree.setProps({ number: 15 });
+
+        expect(tree.text()).to.equal('It’s number 15');
+      });
+
+      it('<DumbSFCMemo /> - should only re-render when number is between 10 and 20', () => {
+        const tree = Wrap(<DumbSFCMemo number={5} />);
+
+        expect(tree.text()).to.equal('It’s number 5');
+
+        tree.setProps({ number: 15 });
+
+        expect(tree.text()).to.equal('It’s number 5');
+      });
+
+      it('<DumbMemoNoCompare /> - should only re-render when prop identity changes', () => {
+        const obj = { toString() { return 5; } };
+
+        const tree = Wrap(<DumbMemoNoCompare number={obj} />);
+
+        expect(tree.text()).to.equal('It’s number 5');
+
+        obj.toString = () => 15;
+        tree.setProps({ number: obj });
+
+        expect(tree.text()).to.equal('It’s number 5');
+
+        tree.setProps({ number: { toString() { return 15; } } });
+
+        expect(tree.text()).to.equal('It’s number 15');
+      });
+
+      it('<DumbSFCMemoNoCompare /> - should only re-render when prop identity changes', () => {
+        const obj = { toString() { return 5; } };
+
+        const tree = Wrap(<DumbSFCMemoNoCompare number={obj} />);
+
+        expect(tree.text()).to.equal('It’s number 5');
+
+        obj.toString = () => 15;
+        tree.setProps({ number: obj });
+
+        expect(tree.text()).to.equal('It’s number 5');
+
+        tree.setProps({ number: { toString() { return 15; } } });
+
+        expect(tree.text()).to.equal('It’s number 15');
+      });
+    });
   });
 }