Support React.memo in ReactShallowRenderer (#14816)

aweary · gaearon · commit b283d75c17ff · 2019-03-15T22:17:09.000Z
* Support React.memo in ReactShallowRenderer

ReactShallowRenderer uses element.type frequently, but with React.memo
elements the actual type is element.type.type. This updates
ReactShallowRenderer so it uses the correct element type for Memo
components and also validates the inner props for the wrapped
components.

* Allow Rect.memo to prevent re-renders

* Support memo(forwardRef())

* Dont call memo comparison function on initial render

* Fix test

* Small tweaks
diff --git a/packages/react-test-renderer/src/ReactShallowRenderer.js b/packages/react-test-renderer/src/ReactShallowRenderer.js
@@ -8,7 +8,7 @@
  */
 
 import React from 'react';
-import {isForwardRef} from 'react-is';
+import {isForwardRef, isMemo, ForwardRef} from 'react-is';
 import describeComponentFrame from 'shared/describeComponentFrame';
 import getComponentName from 'shared/getComponentName';
 import shallowEqual from 'shared/shallowEqual';
@@ -500,7 +500,8 @@ class ReactShallowRenderer {
       element.type,
     );
     invariant(
-      isForwardRef(element) || typeof element.type === 'function',
+      isForwardRef(element) ||
+        (typeof element.type === 'function' || isMemo(element.type)),
       'ReactShallowRenderer render(): Shallow rendering works only with custom ' +
         'components, but the provided element type was `%s`.',
       Array.isArray(element.type)
@@ -514,22 +515,36 @@ class ReactShallowRenderer {
       return;
     }
 
+    const elementType = isMemo(element.type) ? element.type.type : element.type;
+    const previousElement = this._element;
+
     this._rendering = true;
     this._element = element;
-    this._context = getMaskedContext(element.type.contextTypes, context);
+    this._context = getMaskedContext(elementType.contextTypes, context);
+
+    // Inner memo component props aren't currently validated in createElement.
+    if (isMemo(element.type) && elementType.propTypes) {
+      currentlyValidatingElement = element;
+      checkPropTypes(
+        elementType.propTypes,
+        element.props,
+        'prop',
+        getComponentName(elementType),
+        getStackAddendum,
+      );
+    }
 
     if (this._instance) {
-      this._updateClassComponent(element, this._context);
+      this._updateClassComponent(elementType, element, this._context);
     } else {
-      if (shouldConstruct(element.type)) {
-        this._instance = new element.type(
+      if (shouldConstruct(elementType)) {
+        this._instance = new elementType(
           element.props,
           this._context,
           this._updater,
         );
-
-        if (typeof element.type.getDerivedStateFromProps === 'function') {
-          const partialState = element.type.getDerivedStateFromProps.call(
+        if (typeof elementType.getDerivedStateFromProps === 'function') {
+          const partialState = elementType.getDerivedStateFromProps.call(
             null,
             element.props,
             this._instance.state,
@@ -543,39 +558,59 @@ class ReactShallowRenderer {
           }
         }
 
-        if (element.type.hasOwnProperty('contextTypes')) {
+        if (elementType.contextTypes) {
           currentlyValidatingElement = element;
-
           checkPropTypes(
-            element.type.contextTypes,
+            elementType.contextTypes,
             this._context,
             'context',
-            getName(element.type, this._instance),
+            getName(elementType, this._instance),
             getStackAddendum,
           );
 
           currentlyValidatingElement = null;
         }
 
-        this._mountClassComponent(element, this._context);
+        this._mountClassComponent(elementType, element, this._context);
       } else {
-        const prevDispatcher = ReactCurrentDispatcher.current;
-        ReactCurrentDispatcher.current = this._dispatcher;
-        this._prepareToUseHooks(element.type);
-        try {
-          if (isForwardRef(element)) {
-            this._rendered = element.type.render(element.props, element.ref);
-          } else {
-            this._rendered = element.type.call(
-              undefined,
-              element.props,
-              this._context,
-            );
+        let shouldRender = true;
+        if (
+          isMemo(element.type) &&
+          elementType === this._previousComponentIdentity &&
+          previousElement !== null
+        ) {
+          // This is a Memo component that is being re-rendered.
+          const compare = element.type.compare || shallowEqual;
+          if (compare(previousElement.props, element.props)) {
+            shouldRender = false;
+          }
+        }
+        if (shouldRender) {
+          const prevDispatcher = ReactCurrentDispatcher.current;
+          ReactCurrentDispatcher.current = this._dispatcher;
+          this._prepareToUseHooks(elementType);
+          try {
+            // elementType could still be a ForwardRef if it was
+            // nested inside Memo.
+            if (elementType.$$typeof === ForwardRef) {
+              invariant(
+                typeof elementType.render === 'function',
+                'forwardRef requires a render function but was given %s.',
+                typeof elementType.render,
+              );
+              this._rendered = elementType.render.call(
+                undefined,
+                element.props,
+                element.ref,
+              );
+            } else {
+              this._rendered = elementType(element.props, this._context);
+            }
+          } finally {
+            ReactCurrentDispatcher.current = prevDispatcher;
           }
-        } finally {
-          ReactCurrentDispatcher.current = prevDispatcher;
+          this._finishHooks(element, context);
         }
-        this._finishHooks(element, context);
       }
     }
 
@@ -601,7 +636,11 @@ class ReactShallowRenderer {
     this._instance = null;
   }
 
-  _mountClassComponent(element: ReactElement, context: null | Object) {
+  _mountClassComponent(
+    elementType: Function,
+    element: ReactElement,
+    context: null | Object,
+  ) {
     this._instance.context = context;
     this._instance.props = element.props;
     this._instance.state = this._instance.state || null;
@@ -616,7 +655,7 @@ class ReactShallowRenderer {
       // In order to support react-lifecycles-compat polyfilled components,
       // Unsafe lifecycles should not be invoked for components using the new APIs.
       if (
-        typeof element.type.getDerivedStateFromProps !== 'function' &&
+        typeof elementType.getDerivedStateFromProps !== 'function' &&
         typeof this._instance.getSnapshotBeforeUpdate !== 'function'
       ) {
         if (typeof this._instance.componentWillMount === 'function') {
@@ -638,8 +677,12 @@ class ReactShallowRenderer {
     // because DOM refs are not available.
   }
 
-  _updateClassComponent(element: ReactElement, context: null | Object) {
-    const {props, type} = element;
+  _updateClassComponent(
+    elementType: Function,
+    element: ReactElement,
+    context: null | Object,
+  ) {
+    const {props} = element;
 
     const oldState = this._instance.state || emptyObject;
     const oldProps = this._instance.props;
@@ -648,7 +691,7 @@ class ReactShallowRenderer {
       // In order to support react-lifecycles-compat polyfilled components,
       // Unsafe lifecycles should not be invoked for components using the new APIs.
       if (
-        typeof element.type.getDerivedStateFromProps !== 'function' &&
+        typeof elementType.getDerivedStateFromProps !== 'function' &&
         typeof this._instance.getSnapshotBeforeUpdate !== 'function'
       ) {
         if (typeof this._instance.componentWillReceiveProps === 'function') {
@@ -664,8 +707,8 @@ class ReactShallowRenderer {
 
     // Read state after cWRP in case it calls setState
     let state = this._newState || oldState;
-    if (typeof type.getDerivedStateFromProps === 'function') {
-      const partialState = type.getDerivedStateFromProps.call(
+    if (typeof elementType.getDerivedStateFromProps === 'function') {
+      const partialState = elementType.getDerivedStateFromProps.call(
         null,
         props,
         state,
@@ -685,7 +728,10 @@ class ReactShallowRenderer {
         state,
         context,
       );
-    } else if (type.prototype && type.prototype.isPureReactComponent) {
+    } else if (
+      elementType.prototype &&
+      elementType.prototype.isPureReactComponent
+    ) {
       shouldUpdate =
         !shallowEqual(oldProps, props) || !shallowEqual(oldState, state);
     }
@@ -694,7 +740,7 @@ class ReactShallowRenderer {
       // In order to support react-lifecycles-compat polyfilled components,
       // Unsafe lifecycles should not be invoked for components using the new APIs.
       if (
-        typeof element.type.getDerivedStateFromProps !== 'function' &&
+        typeof elementType.getDerivedStateFromProps !== 'function' &&
         typeof this._instance.getSnapshotBeforeUpdate !== 'function'
       ) {
         if (typeof this._instance.componentWillUpdate === 'function') {
@@ -729,7 +775,8 @@ function getDisplayName(element) {
   } else if (typeof element.type === 'string') {
     return element.type;
   } else {
-    return element.type.displayName || element.type.name || 'Unknown';
+    const elementType = isMemo(element.type) ? element.type.type : element.type;
+    return elementType.displayName || elementType.name || 'Unknown';
   }
 }
 
diff --git a/packages/react-test-renderer/src/__tests__/ReactShallowRenderer-test.js b/packages/react-test-renderer/src/__tests__/ReactShallowRenderer-test.js
@@ -1454,4 +1454,115 @@ describe('ReactShallowRenderer', () => {
     shallowRenderer.render(<Foo foo="bar" />);
     expect(logs).toEqual([undefined]);
   });
+
+  it('should handle memo', () => {
+    function Foo() {
+      return <div>foo</div>;
+    }
+    const MemoFoo = React.memo(Foo);
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<MemoFoo />);
+  });
+
+  it('should enable React.memo to prevent a re-render', () => {
+    const logs = [];
+    const Foo = React.memo(({count}) => {
+      logs.push(`Foo: ${count}`);
+      return <div>{count}</div>;
+    });
+    const Bar = React.memo(({count}) => {
+      logs.push(`Bar: ${count}`);
+      return <div>{count}</div>;
+    });
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<Foo count={1} />);
+    expect(logs).toEqual(['Foo: 1']);
+    logs.length = 0;
+    // Rendering the same element with the same props should be prevented
+    shallowRenderer.render(<Foo count={1} />);
+    expect(logs).toEqual([]);
+    // A different element with the same props should cause a re-render
+    shallowRenderer.render(<Bar count={1} />);
+    expect(logs).toEqual(['Bar: 1']);
+  });
+
+  it('should respect a custom comparison function with React.memo', () => {
+    let renderCount = 0;
+    function areEqual(props, nextProps) {
+      return props.foo === nextProps.foo;
+    }
+    const Foo = React.memo(({foo, bar}) => {
+      renderCount++;
+      return (
+        <div>
+          {foo} {bar}
+        </div>
+      );
+    }, areEqual);
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<Foo foo={1} bar={1} />);
+    expect(renderCount).toBe(1);
+    // Change a prop that the comparison funciton ignores
+    shallowRenderer.render(<Foo foo={1} bar={2} />);
+    expect(renderCount).toBe(1);
+    shallowRenderer.render(<Foo foo={2} bar={2} />);
+    expect(renderCount).toBe(2);
+  });
+
+  it('should not call the comparison function with React.memo on the initial render', () => {
+    const areEqual = jest.fn(() => false);
+    const SomeComponent = React.memo(({foo}) => {
+      return <div>{foo}</div>;
+    }, areEqual);
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SomeComponent foo={1} />);
+    expect(areEqual).not.toHaveBeenCalled();
+    expect(shallowRenderer.getRenderOutput()).toEqual(<div>{1}</div>);
+  });
+
+  it('should handle memo(forwardRef())', () => {
+    const testRef = React.createRef();
+    const SomeComponent = React.forwardRef((props, ref) => {
+      expect(ref).toEqual(testRef);
+      return (
+        <div>
+          <span className="child1" />
+          <span className="child2" />
+        </div>
+      );
+    });
+
+    const SomeMemoComponent = React.memo(SomeComponent);
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SomeMemoComponent ref={testRef} />);
+
+    expect(result.type).toBe('div');
+    expect(result.props.children).toEqual([
+      <span className="child1" />,
+      <span className="child2" />,
+    ]);
+  });
+
+  it('should warn for forwardRef(memo())', () => {
+    const testRef = React.createRef();
+    const SomeMemoComponent = React.memo(({foo}) => {
+      return <div>{foo}</div>;
+    });
+    const shallowRenderer = createRenderer();
+    expect(() => {
+      expect(() => {
+        const SomeComponent = React.forwardRef(SomeMemoComponent);
+        shallowRenderer.render(<SomeComponent ref={testRef} />);
+      }).toWarnDev(
+        'Warning: forwardRef requires a render function but received ' +
+          'a `memo` component. Instead of forwardRef(memo(...)), use ' +
+          'memo(forwardRef(...))',
+        {withoutStack: true},
+      );
+    }).toThrowError(
+      'forwardRef requires a render function but was given object.',
+    );
+  });
 });
diff --git a/packages/react-test-renderer/src/__tests__/ReactShallowRendererMemo-test.js b/packages/react-test-renderer/src/__tests__/ReactShallowRendererMemo-test.js
