From b283d75c17ff40cba1a49956d23b3985b9eb6abe Mon Sep 17 00:00:00 2001
From: Brandon Dail <aweary@users.noreply.github.com>
Date: Fri, 15 Mar 2019 15:17:09 -0700
Subject: [PATCH] Support React.memo in ReactShallowRenderer (#14816)

* 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
---
 .../src/ReactShallowRenderer.js               |  125 +-
 .../__tests__/ReactShallowRenderer-test.js    |  111 ++
 .../ReactShallowRendererMemo-test.js          | 1520 +++++++++++++++++
 3 files changed, 1717 insertions(+), 39 deletions(-)
 create mode 100644 packages/react-test-renderer/src/__tests__/ReactShallowRendererMemo-test.js

diff --git a/packages/react-test-renderer/src/ReactShallowRenderer.js b/packages/react-test-renderer/src/ReactShallowRenderer.js
index a2775eb2c446c..95081009a0ea5 100644
--- 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
index d83dada160f90..df7a08d9d5c0d 100644
--- 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
new file mode 100644
index 0000000000000..839524661459f
--- /dev/null
+++ b/packages/react-test-renderer/src/__tests__/ReactShallowRendererMemo-test.js
@@ -0,0 +1,1520 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails react-core
+ * @jest-environment node
+ */
+
+'use strict';
+
+let createRenderer;
+let PropTypes;
+let React;
+
+describe('ReactShallowRendererMemo', () => {
+  beforeEach(() => {
+    jest.resetModules();
+
+    createRenderer = require('react-test-renderer/shallow').createRenderer;
+    PropTypes = require('prop-types');
+    React = require('react');
+  });
+
+  it('should call all of the legacy lifecycle hooks', () => {
+    const logs = [];
+    const logger = message => () => logs.push(message) || true;
+
+    const SomeComponent = React.memo(
+      class SomeComponent extends React.Component {
+        UNSAFE_componentWillMount = logger('componentWillMount');
+        componentDidMount = logger('componentDidMount');
+        UNSAFE_componentWillReceiveProps = logger('componentWillReceiveProps');
+        shouldComponentUpdate = logger('shouldComponentUpdate');
+        UNSAFE_componentWillUpdate = logger('componentWillUpdate');
+        componentDidUpdate = logger('componentDidUpdate');
+        componentWillUnmount = logger('componentWillUnmount');
+        render() {
+          return <div />;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SomeComponent foo={1} />);
+
+    // Calling cDU might lead to problems with host component references.
+    // Since our components aren't really mounted, refs won't be available.
+    expect(logs).toEqual(['componentWillMount']);
+
+    logs.splice(0);
+
+    const instance = shallowRenderer.getMountedInstance();
+    instance.setState({});
+
+    expect(logs).toEqual(['shouldComponentUpdate', 'componentWillUpdate']);
+
+    logs.splice(0);
+
+    shallowRenderer.render(<SomeComponent foo={2} />);
+
+    // The previous shallow renderer did not trigger cDU for props changes.
+    expect(logs).toEqual([
+      'componentWillReceiveProps',
+      'shouldComponentUpdate',
+      'componentWillUpdate',
+    ]);
+  });
+
+  it('should call all of the new lifecycle hooks', () => {
+    const logs = [];
+    const logger = message => () => logs.push(message) || true;
+
+    const SomeComponent = React.memo(
+      class SomeComponent extends React.Component {
+        state = {};
+        static getDerivedStateFromProps = logger('getDerivedStateFromProps');
+        componentDidMount = logger('componentDidMount');
+        shouldComponentUpdate = logger('shouldComponentUpdate');
+        componentDidUpdate = logger('componentDidUpdate');
+        componentWillUnmount = logger('componentWillUnmount');
+        render() {
+          return <div />;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SomeComponent foo={1} />);
+
+    // Calling cDU might lead to problems with host component references.
+    // Since our components aren't really mounted, refs won't be available.
+    expect(logs).toEqual(['getDerivedStateFromProps']);
+
+    logs.splice(0);
+
+    const instance = shallowRenderer.getMountedInstance();
+    instance.setState({});
+
+    expect(logs).toEqual(['getDerivedStateFromProps', 'shouldComponentUpdate']);
+
+    logs.splice(0);
+
+    shallowRenderer.render(<SomeComponent foo={2} />);
+
+    // The previous shallow renderer did not trigger cDU for props changes.
+    expect(logs).toEqual(['getDerivedStateFromProps', 'shouldComponentUpdate']);
+  });
+
+  it('should not invoke deprecated lifecycles (cWM/cWRP/cWU) if new static gDSFP is present', () => {
+    const Component = React.memo(
+      class Component extends React.Component {
+        state = {};
+        static getDerivedStateFromProps() {
+          return null;
+        }
+        componentWillMount() {
+          throw Error('unexpected');
+        }
+        componentWillReceiveProps() {
+          throw Error('unexpected');
+        }
+        componentWillUpdate() {
+          throw Error('unexpected');
+        }
+        render() {
+          return null;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<Component />);
+  });
+
+  it('should not invoke deprecated lifecycles (cWM/cWRP/cWU) if new getSnapshotBeforeUpdate is present', () => {
+    const Component = React.memo(
+      class Component extends React.Component {
+        getSnapshotBeforeUpdate() {
+          return null;
+        }
+        componentWillMount() {
+          throw Error('unexpected');
+        }
+        componentWillReceiveProps() {
+          throw Error('unexpected');
+        }
+        componentWillUpdate() {
+          throw Error('unexpected');
+        }
+        render() {
+          return null;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<Component value={1} />);
+    shallowRenderer.render(<Component value={2} />);
+  });
+
+  it('should not call getSnapshotBeforeUpdate or componentDidUpdate when updating since refs wont exist', () => {
+    const Component = React.memo(
+      class Component extends React.Component {
+        getSnapshotBeforeUpdate() {
+          throw Error('unexpected');
+        }
+        componentDidUpdate() {
+          throw Error('unexpected');
+        }
+        render() {
+          return null;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<Component value={1} />);
+    shallowRenderer.render(<Component value={2} />);
+  });
+
+  it('should only render 1 level deep', () => {
+    const Parent = React.memo(function Parent() {
+      return (
+        <div>
+          <Child />
+        </div>
+      );
+    });
+
+    function Child() {
+      throw Error('This component should not render');
+    }
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(React.createElement(Parent));
+  });
+
+  it('should have shallow rendering', () => {
+    const SomeComponent = React.memo(
+      class SomeComponent extends React.Component {
+        render() {
+          return (
+            <div>
+              <span className="child1" />
+              <span className="child2" />
+            </div>
+          );
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SomeComponent />);
+
+    expect(result.type).toBe('div');
+    expect(result.props.children).toEqual([
+      <span className="child1" />,
+      <span className="child2" />,
+    ]);
+  });
+
+  it('should handle Profiler', () => {
+    const SomeComponent = React.memo(
+      class SomeComponent extends React.Component {
+        render() {
+          return (
+            <React.unstable_Profiler id="test" onRender={jest.fn()}>
+              <div>
+                <span className="child1" />
+                <span className="child2" />
+              </div>
+            </React.unstable_Profiler>
+          );
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SomeComponent />);
+
+    expect(result.type).toBe(React.unstable_Profiler);
+    expect(result.props.children).toEqual(
+      <div>
+        <span className="child1" />
+        <span className="child2" />
+      </div>,
+    );
+  });
+
+  it('should enable shouldComponentUpdate to prevent a re-render', () => {
+    let renderCounter = 0;
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {update: false};
+        shouldComponentUpdate(nextProps, nextState) {
+          return this.state.update !== nextState.update;
+        }
+        render() {
+          renderCounter++;
+          return <div>{`${renderCounter}`}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SimpleComponent />);
+    expect(shallowRenderer.getRenderOutput()).toEqual(<div>1</div>);
+
+    const instance = shallowRenderer.getMountedInstance();
+    instance.setState({update: false});
+    expect(shallowRenderer.getRenderOutput()).toEqual(<div>1</div>);
+
+    instance.setState({update: true});
+    expect(shallowRenderer.getRenderOutput()).toEqual(<div>2</div>);
+  });
+
+  it('should enable PureComponent to prevent a re-render', () => {
+    let renderCounter = 0;
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.PureComponent {
+        state = {update: false};
+        render() {
+          renderCounter++;
+          return <div>{`${renderCounter}`}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SimpleComponent />);
+    expect(shallowRenderer.getRenderOutput()).toEqual(<div>1</div>);
+
+    const instance = shallowRenderer.getMountedInstance();
+    instance.setState({update: false});
+    expect(shallowRenderer.getRenderOutput()).toEqual(<div>1</div>);
+
+    instance.setState({update: true});
+    expect(shallowRenderer.getRenderOutput()).toEqual(<div>2</div>);
+  });
+
+  it('should not run shouldComponentUpdate during forced update', () => {
+    let scuCounter = 0;
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {count: 1};
+        shouldComponentUpdate() {
+          scuCounter++;
+          return false;
+        }
+        render() {
+          return <div>{`${this.state.count}`}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SimpleComponent />);
+    expect(scuCounter).toEqual(0);
+    expect(shallowRenderer.getRenderOutput()).toEqual(<div>1</div>);
+
+    // Force update the initial state. sCU should not fire.
+    const instance = shallowRenderer.getMountedInstance();
+    instance.forceUpdate();
+    expect(scuCounter).toEqual(0);
+    expect(shallowRenderer.getRenderOutput()).toEqual(<div>1</div>);
+
+    // Setting state updates the instance, but doesn't re-render
+    // because sCU returned false.
+    instance.setState(state => ({count: state.count + 1}));
+    expect(scuCounter).toEqual(1);
+    expect(instance.state.count).toEqual(2);
+    expect(shallowRenderer.getRenderOutput()).toEqual(<div>1</div>);
+
+    // A force update updates the render output, but doesn't call sCU.
+    instance.forceUpdate();
+    expect(scuCounter).toEqual(1);
+    expect(instance.state.count).toEqual(2);
+    expect(shallowRenderer.getRenderOutput()).toEqual(<div>2</div>);
+  });
+
+  it('should rerender when calling forceUpdate', () => {
+    let renderCounter = 0;
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        render() {
+          renderCounter += 1;
+          return <div />;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SimpleComponent />);
+    expect(renderCounter).toEqual(1);
+
+    const instance = shallowRenderer.getMountedInstance();
+    instance.forceUpdate();
+    expect(renderCounter).toEqual(2);
+  });
+
+  it('should shallow render a function component', () => {
+    function SomeComponent(props, context) {
+      return (
+        <div>
+          <div>{props.foo}</div>
+          <div>{context.bar}</div>
+          <span className="child1" />
+          <span className="child2" />
+        </div>
+      );
+    }
+    const SomeMemoComponent = React.memo(SomeComponent);
+
+    SomeComponent.contextTypes = {
+      bar: PropTypes.string,
+    };
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SomeMemoComponent foo={'FOO'} />, {
+      bar: 'BAR',
+    });
+
+    expect(result.type).toBe('div');
+    expect(result.props.children).toEqual([
+      <div>FOO</div>,
+      <div>BAR</div>,
+      <span className="child1" />,
+      <span className="child2" />,
+    ]);
+  });
+
+  it('should shallow render a component returning strings directly from render', () => {
+    const Text = React.memo(({value}) => value);
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<Text value="foo" />);
+    expect(result).toEqual('foo');
+  });
+
+  it('should shallow render a component returning numbers directly from render', () => {
+    const Text = React.memo(({value}) => value);
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<Text value={10} />);
+    expect(result).toEqual(10);
+  });
+
+  it('should shallow render a fragment', () => {
+    class SomeComponent extends React.Component {
+      render() {
+        return <div />;
+      }
+    }
+    class Fragment extends React.Component {
+      render() {
+        return [<div key="a" />, <span key="b" />, <SomeComponent />];
+      }
+    }
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<Fragment />);
+    expect(result).toEqual([
+      <div key="a" />,
+      <span key="b" />,
+      <SomeComponent />,
+    ]);
+  });
+
+  it('should shallow render a React.fragment', () => {
+    class SomeComponent extends React.Component {
+      render() {
+        return <div />;
+      }
+    }
+    class Fragment extends React.Component {
+      render() {
+        return (
+          <React.Fragment>
+            <div />
+            <span />
+            <SomeComponent />
+          </React.Fragment>
+        );
+      }
+    }
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<Fragment />);
+    expect(result).toEqual(
+      <React.Fragment>
+        <div />
+        <span />
+        <SomeComponent />
+      </React.Fragment>,
+    );
+  });
+
+  it('should throw for invalid elements', () => {
+    class SomeComponent extends React.Component {
+      render() {
+        return <div />;
+      }
+    }
+
+    const shallowRenderer = createRenderer();
+    expect(() => shallowRenderer.render(SomeComponent)).toThrowError(
+      'ReactShallowRenderer render(): Invalid component element. Instead of ' +
+        'passing a component class, make sure to instantiate it by passing it ' +
+        'to React.createElement.',
+    );
+    expect(() => shallowRenderer.render(<div />)).toThrowError(
+      'ReactShallowRenderer render(): Shallow rendering works only with ' +
+        'custom components, not primitives (div). Instead of calling ' +
+        '`.render(el)` and inspecting the rendered output, look at `el.props` ' +
+        'directly instead.',
+    );
+  });
+
+  it('should have shallow unmounting', () => {
+    const componentWillUnmount = jest.fn();
+
+    class SomeComponent extends React.Component {
+      componentWillUnmount = componentWillUnmount;
+      render() {
+        return <div />;
+      }
+    }
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SomeComponent />);
+    shallowRenderer.unmount();
+
+    expect(componentWillUnmount).toBeCalled();
+  });
+
+  it('can shallow render to null', () => {
+    class SomeComponent extends React.Component {
+      render() {
+        return null;
+      }
+    }
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SomeComponent />);
+
+    expect(result).toBe(null);
+  });
+
+  it('can shallow render with a ref', () => {
+    class SomeComponent extends React.Component {
+      render() {
+        return <div ref="hello" />;
+      }
+    }
+
+    const shallowRenderer = createRenderer();
+    // Shouldn't crash.
+    shallowRenderer.render(<SomeComponent />);
+  });
+
+  it('lets you update shallowly rendered components', () => {
+    class SomeComponent extends React.Component {
+      state = {clicked: false};
+
+      onClick = () => {
+        this.setState({clicked: true});
+      };
+
+      render() {
+        const className = this.state.clicked ? 'was-clicked' : '';
+
+        if (this.props.aNew === 'prop') {
+          return (
+            <a href="#" onClick={this.onClick} className={className}>
+              Test link
+            </a>
+          );
+        } else {
+          return (
+            <div>
+              <span className="child1" />
+              <span className="child2" />
+            </div>
+          );
+        }
+      }
+    }
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SomeComponent />);
+    expect(result.type).toBe('div');
+    expect(result.props.children).toEqual([
+      <span className="child1" />,
+      <span className="child2" />,
+    ]);
+
+    const updatedResult = shallowRenderer.render(<SomeComponent aNew="prop" />);
+    expect(updatedResult.type).toBe('a');
+
+    const mockEvent = {};
+    updatedResult.props.onClick(mockEvent);
+
+    const updatedResultCausedByClick = shallowRenderer.getRenderOutput();
+    expect(updatedResultCausedByClick.type).toBe('a');
+    expect(updatedResultCausedByClick.props.className).toBe('was-clicked');
+  });
+
+  it('can access the mounted component instance', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        someMethod = () => {
+          return this.props.n;
+        };
+
+        render() {
+          return <div>{this.props.n}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SimpleComponent n={5} />);
+    expect(shallowRenderer.getMountedInstance().someMethod()).toEqual(5);
+  });
+
+  it('can shallowly render components with contextTypes', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        static contextTypes = {
+          name: PropTypes.string,
+        };
+
+        render() {
+          return <div />;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SimpleComponent />);
+    expect(result).toEqual(<div />);
+  });
+
+  it('passes expected params to legacy component lifecycle methods', () => {
+    const componentDidUpdateParams = [];
+    const componentWillReceivePropsParams = [];
+    const componentWillUpdateParams = [];
+    const setStateParams = [];
+    const shouldComponentUpdateParams = [];
+
+    const initialProp = {prop: 'init prop'};
+    const initialState = {state: 'init state'};
+    const initialContext = {context: 'init context'};
+    const updatedState = {state: 'updated state'};
+    const updatedProp = {prop: 'updated prop'};
+    const updatedContext = {context: 'updated context'};
+
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        constructor(props, context) {
+          super(props, context);
+          this.state = initialState;
+        }
+        static contextTypes = {
+          context: PropTypes.string,
+        };
+        componentDidUpdate(...args) {
+          componentDidUpdateParams.push(...args);
+        }
+        UNSAFE_componentWillReceiveProps(...args) {
+          componentWillReceivePropsParams.push(...args);
+          this.setState((...innerArgs) => {
+            setStateParams.push(...innerArgs);
+            return updatedState;
+          });
+        }
+        UNSAFE_componentWillUpdate(...args) {
+          componentWillUpdateParams.push(...args);
+        }
+        shouldComponentUpdate(...args) {
+          shouldComponentUpdateParams.push(...args);
+          return true;
+        }
+        render() {
+          return null;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(
+      React.createElement(SimpleComponent, initialProp),
+      initialContext,
+    );
+    expect(componentDidUpdateParams).toEqual([]);
+    expect(componentWillReceivePropsParams).toEqual([]);
+    expect(componentWillUpdateParams).toEqual([]);
+    expect(setStateParams).toEqual([]);
+    expect(shouldComponentUpdateParams).toEqual([]);
+
+    // Lifecycle hooks should be invoked with the correct prev/next params on update.
+    shallowRenderer.render(
+      React.createElement(SimpleComponent, updatedProp),
+      updatedContext,
+    );
+    expect(componentWillReceivePropsParams).toEqual([
+      updatedProp,
+      updatedContext,
+    ]);
+    expect(setStateParams).toEqual([initialState, initialProp]);
+    expect(shouldComponentUpdateParams).toEqual([
+      updatedProp,
+      updatedState,
+      updatedContext,
+    ]);
+    expect(componentWillUpdateParams).toEqual([
+      updatedProp,
+      updatedState,
+      updatedContext,
+    ]);
+    expect(componentDidUpdateParams).toEqual([]);
+  });
+
+  it('passes expected params to new component lifecycle methods', () => {
+    const componentDidUpdateParams = [];
+    const getDerivedStateFromPropsParams = [];
+    const shouldComponentUpdateParams = [];
+
+    const initialProp = {prop: 'init prop'};
+    const initialState = {state: 'init state'};
+    const initialContext = {context: 'init context'};
+    const updatedProp = {prop: 'updated prop'};
+    const updatedContext = {context: 'updated context'};
+
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        constructor(props, context) {
+          super(props, context);
+          this.state = initialState;
+        }
+        static contextTypes = {
+          context: PropTypes.string,
+        };
+        componentDidUpdate(...args) {
+          componentDidUpdateParams.push(...args);
+        }
+        static getDerivedStateFromProps(...args) {
+          getDerivedStateFromPropsParams.push(args);
+          return null;
+        }
+        shouldComponentUpdate(...args) {
+          shouldComponentUpdateParams.push(...args);
+          return true;
+        }
+        render() {
+          return null;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+
+    // The only lifecycle hook that should be invoked on initial render
+    // Is the static getDerivedStateFromProps() methods
+    shallowRenderer.render(
+      React.createElement(SimpleComponent, initialProp),
+      initialContext,
+    );
+    expect(getDerivedStateFromPropsParams).toEqual([
+      [initialProp, initialState],
+    ]);
+    expect(componentDidUpdateParams).toEqual([]);
+    expect(shouldComponentUpdateParams).toEqual([]);
+
+    // Lifecycle hooks should be invoked with the correct prev/next params on update.
+    shallowRenderer.render(
+      React.createElement(SimpleComponent, updatedProp),
+      updatedContext,
+    );
+    expect(getDerivedStateFromPropsParams).toEqual([
+      [initialProp, initialState],
+      [updatedProp, initialState],
+    ]);
+    expect(shouldComponentUpdateParams).toEqual([
+      updatedProp,
+      initialState,
+      updatedContext,
+    ]);
+    expect(componentDidUpdateParams).toEqual([]);
+  });
+
+  it('can shallowly render components with ref as function', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {clicked: false};
+
+        handleUserClick = () => {
+          this.setState({clicked: true});
+        };
+
+        render() {
+          return (
+            <div
+              ref={() => {}}
+              onClick={this.handleUserClick}
+              className={this.state.clicked ? 'clicked' : ''}
+            />
+          );
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SimpleComponent />);
+    let result = shallowRenderer.getRenderOutput();
+    expect(result.type).toEqual('div');
+    expect(result.props.className).toEqual('');
+    result.props.onClick();
+
+    result = shallowRenderer.getRenderOutput();
+    expect(result.type).toEqual('div');
+    expect(result.props.className).toEqual('clicked');
+  });
+
+  it('can initialize state via static getDerivedStateFromProps', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {
+          count: 1,
+        };
+
+        static getDerivedStateFromProps(props, prevState) {
+          return {
+            count: prevState.count + props.incrementBy,
+            other: 'foobar',
+          };
+        }
+
+        render() {
+          return (
+            <div>{`count:${this.state.count}, other:${this.state.other}`}</div>
+          );
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SimpleComponent incrementBy={2} />);
+    expect(result).toEqual(<div>count:3, other:foobar</div>);
+  });
+
+  it('can setState in componentWillMount when shallow rendering', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        UNSAFE_componentWillMount() {
+          this.setState({groovy: 'doovy'});
+        }
+
+        render() {
+          return <div>{this.state.groovy}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SimpleComponent />);
+    expect(result).toEqual(<div>doovy</div>);
+  });
+
+  it('can setState in componentWillMount repeatedly when shallow rendering', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {
+          separator: '-',
+        };
+
+        UNSAFE_componentWillMount() {
+          this.setState({groovy: 'doovy'});
+          this.setState({doovy: 'groovy'});
+        }
+
+        render() {
+          const {groovy, doovy, separator} = this.state;
+
+          return <div>{`${groovy}${separator}${doovy}`}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SimpleComponent />);
+    expect(result).toEqual(<div>doovy-groovy</div>);
+  });
+
+  it('can setState in componentWillMount with an updater function repeatedly when shallow rendering', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {
+          separator: '-',
+        };
+
+        UNSAFE_componentWillMount() {
+          this.setState(state => ({groovy: 'doovy'}));
+          this.setState(state => ({doovy: state.groovy}));
+        }
+
+        render() {
+          const {groovy, doovy, separator} = this.state;
+
+          return <div>{`${groovy}${separator}${doovy}`}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SimpleComponent />);
+    expect(result).toEqual(<div>doovy-doovy</div>);
+  });
+
+  it('can setState in componentWillReceiveProps when shallow rendering', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {count: 0};
+
+        UNSAFE_componentWillReceiveProps(nextProps) {
+          if (nextProps.updateState) {
+            this.setState({count: 1});
+          }
+        }
+
+        render() {
+          return <div>{this.state.count}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    let result = shallowRenderer.render(
+      <SimpleComponent updateState={false} />,
+    );
+    expect(result.props.children).toEqual(0);
+
+    result = shallowRenderer.render(<SimpleComponent updateState={true} />);
+    expect(result.props.children).toEqual(1);
+  });
+
+  it('can update state with static getDerivedStateFromProps when shallow rendering', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {count: 1};
+
+        static getDerivedStateFromProps(nextProps, prevState) {
+          if (nextProps.updateState) {
+            return {count: nextProps.incrementBy + prevState.count};
+          }
+
+          return null;
+        }
+
+        render() {
+          return <div>{this.state.count}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    let result = shallowRenderer.render(
+      <SimpleComponent updateState={false} incrementBy={0} />,
+    );
+    expect(result.props.children).toEqual(1);
+
+    result = shallowRenderer.render(
+      <SimpleComponent updateState={true} incrementBy={2} />,
+    );
+    expect(result.props.children).toEqual(3);
+
+    result = shallowRenderer.render(
+      <SimpleComponent updateState={false} incrementBy={2} />,
+    );
+    expect(result.props.children).toEqual(3);
+  });
+
+  it('should not override state with stale values if prevState is spread within getDerivedStateFromProps', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {value: 0};
+
+        static getDerivedStateFromProps(nextProps, prevState) {
+          return {...prevState};
+        }
+
+        updateState = () => {
+          this.setState(state => ({value: state.value + 1}));
+        };
+
+        render() {
+          return <div>{`value:${this.state.value}`}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    let result = shallowRenderer.render(<SimpleComponent />);
+    expect(result).toEqual(<div>value:0</div>);
+
+    let instance = shallowRenderer.getMountedInstance();
+    instance.updateState();
+    result = shallowRenderer.getRenderOutput();
+    expect(result).toEqual(<div>value:1</div>);
+  });
+
+  it('should pass previous state to shouldComponentUpdate even with getDerivedStateFromProps', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = {
+            value: props.value,
+          };
+        }
+
+        static getDerivedStateFromProps(nextProps, prevState) {
+          if (nextProps.value === prevState.value) {
+            return null;
+          }
+          return {value: nextProps.value};
+        }
+
+        shouldComponentUpdate(nextProps, nextState) {
+          return nextState.value !== this.state.value;
+        }
+
+        render() {
+          return <div>{`value:${this.state.value}`}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const initialResult = shallowRenderer.render(
+      <SimpleComponent value="initial" />,
+    );
+    expect(initialResult).toEqual(<div>value:initial</div>);
+    const updatedResult = shallowRenderer.render(
+      <SimpleComponent value="updated" />,
+    );
+    expect(updatedResult).toEqual(<div>value:updated</div>);
+  });
+
+  it('can setState with an updater function', () => {
+    let instance;
+
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {
+          counter: 0,
+        };
+
+        render() {
+          instance = this;
+          return (
+            <button ref="button" onClick={this.onClick}>
+              {this.state.counter}
+            </button>
+          );
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    let result = shallowRenderer.render(<SimpleComponent defaultCount={1} />);
+    expect(result.props.children).toEqual(0);
+
+    instance.setState((state, props) => {
+      return {counter: props.defaultCount + 1};
+    });
+
+    result = shallowRenderer.getRenderOutput();
+    expect(result.props.children).toEqual(2);
+  });
+
+  it('can access component instance from setState updater function', done => {
+    let instance;
+
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {};
+
+        render() {
+          instance = this;
+          return null;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SimpleComponent />);
+
+    instance.setState(function updater(state, props) {
+      expect(this).toBe(instance);
+      done();
+    });
+  });
+
+  it('can setState with a callback', () => {
+    let instance;
+
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {
+          counter: 0,
+        };
+        render() {
+          instance = this;
+          return <p>{this.state.counter}</p>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SimpleComponent />);
+    expect(result.props.children).toBe(0);
+
+    const callback = jest.fn(function() {
+      expect(this).toBe(instance);
+    });
+
+    instance.setState({counter: 1}, callback);
+
+    const updated = shallowRenderer.getRenderOutput();
+    expect(updated.props.children).toBe(1);
+    expect(callback).toHaveBeenCalled();
+  });
+
+  it('can replaceState with a callback', () => {
+    let instance;
+
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {
+          counter: 0,
+        };
+        render() {
+          instance = this;
+          return <p>{this.state.counter}</p>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SimpleComponent />);
+    expect(result.props.children).toBe(0);
+
+    const callback = jest.fn(function() {
+      expect(this).toBe(instance);
+    });
+
+    // No longer a public API, but we can test that it works internally by
+    // reaching into the updater.
+    shallowRenderer._updater.enqueueReplaceState(
+      instance,
+      {counter: 1},
+      callback,
+    );
+
+    const updated = shallowRenderer.getRenderOutput();
+    expect(updated.props.children).toBe(1);
+    expect(callback).toHaveBeenCalled();
+  });
+
+  it('can forceUpdate with a callback', () => {
+    let instance;
+
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        state = {
+          counter: 0,
+        };
+        render() {
+          instance = this;
+          return <p>{this.state.counter}</p>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SimpleComponent />);
+    expect(result.props.children).toBe(0);
+
+    const callback = jest.fn(function() {
+      expect(this).toBe(instance);
+    });
+
+    instance.forceUpdate(callback);
+
+    const updated = shallowRenderer.getRenderOutput();
+    expect(updated.props.children).toBe(0);
+    expect(callback).toHaveBeenCalled();
+  });
+
+  it('can pass context when shallowly rendering', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        static contextTypes = {
+          name: PropTypes.string,
+        };
+
+        render() {
+          return <div>{this.context.name}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SimpleComponent />, {
+      name: 'foo',
+    });
+    expect(result).toEqual(<div>foo</div>);
+  });
+
+  it('should track context across updates', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        static contextTypes = {
+          foo: PropTypes.string,
+        };
+
+        state = {
+          bar: 'bar',
+        };
+
+        render() {
+          return <div>{`${this.context.foo}:${this.state.bar}`}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    let result = shallowRenderer.render(<SimpleComponent />, {
+      foo: 'foo',
+    });
+    expect(result).toEqual(<div>foo:bar</div>);
+
+    const instance = shallowRenderer.getMountedInstance();
+    instance.setState({bar: 'baz'});
+
+    result = shallowRenderer.getRenderOutput();
+    expect(result).toEqual(<div>foo:baz</div>);
+  });
+
+  it('should filter context by contextTypes', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        static contextTypes = {
+          foo: PropTypes.string,
+        };
+        render() {
+          return <div>{`${this.context.foo}:${this.context.bar}`}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    let result = shallowRenderer.render(<SimpleComponent />, {
+      foo: 'foo',
+      bar: 'bar',
+    });
+    expect(result).toEqual(<div>foo:undefined</div>);
+  });
+
+  it('can fail context when shallowly rendering', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        static contextTypes = {
+          name: PropTypes.string.isRequired,
+        };
+
+        render() {
+          return <div>{this.context.name}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    expect(() => shallowRenderer.render(<SimpleComponent />)).toWarnDev(
+      'Warning: Failed context type: The context `name` is marked as ' +
+        'required in `SimpleComponent`, but its value is `undefined`.\n' +
+        '    in SimpleComponent (at **)',
+    );
+  });
+
+  it('should warn about propTypes (but only once)', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        static propTypes = {
+          name: PropTypes.string.isRequired,
+        };
+
+        render() {
+          return React.createElement('div', null, this.props.name);
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    expect(() =>
+      shallowRenderer.render(React.createElement(SimpleComponent, {name: 123})),
+    ).toWarnDev(
+      'Warning: Failed prop type: Invalid prop `name` of type `number` ' +
+        'supplied to `SimpleComponent`, expected `string`.\n' +
+        '    in SimpleComponent',
+    );
+  });
+
+  it('should enable rendering of cloned element', () => {
+    const SimpleComponent = React.memo(
+      class SimpleComponent extends React.Component {
+        constructor(props) {
+          super(props);
+
+          this.state = {
+            bar: 'bar',
+          };
+        }
+
+        render() {
+          return <div>{`${this.props.foo}:${this.state.bar}`}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    const el = <SimpleComponent foo="foo" />;
+    let result = shallowRenderer.render(el);
+    expect(result).toEqual(<div>foo:bar</div>);
+
+    const cloned = React.cloneElement(el, {foo: 'baz'});
+    result = shallowRenderer.render(cloned);
+    expect(result).toEqual(<div>baz:bar</div>);
+  });
+
+  it('this.state should be updated on setState callback inside componentWillMount', () => {
+    let stateSuccessfullyUpdated = false;
+
+    const Component = React.memo(
+      class Component extends React.Component {
+        constructor(props, context) {
+          super(props, context);
+          this.state = {
+            hasUpdatedState: false,
+          };
+        }
+
+        UNSAFE_componentWillMount() {
+          this.setState(
+            {hasUpdatedState: true},
+            () => (stateSuccessfullyUpdated = this.state.hasUpdatedState),
+          );
+        }
+
+        render() {
+          return <div>{this.props.children}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<Component />);
+    expect(stateSuccessfullyUpdated).toBe(true);
+  });
+
+  it('should handle multiple callbacks', () => {
+    const mockFn = jest.fn();
+    const shallowRenderer = createRenderer();
+
+    const Component = React.memo(
+      class Component extends React.Component {
+        constructor(props, context) {
+          super(props, context);
+          this.state = {
+            foo: 'foo',
+          };
+        }
+
+        UNSAFE_componentWillMount() {
+          this.setState({foo: 'bar'}, () => mockFn());
+          this.setState({foo: 'foobar'}, () => mockFn());
+        }
+
+        render() {
+          return <div>{this.state.foo}</div>;
+        }
+      },
+    );
+
+    shallowRenderer.render(<Component />);
+
+    expect(mockFn).toHaveBeenCalledTimes(2);
+
+    // Ensure the callback queue is cleared after the callbacks are invoked
+    const mountedInstance = shallowRenderer.getMountedInstance();
+    mountedInstance.setState({foo: 'bar'}, () => mockFn());
+    expect(mockFn).toHaveBeenCalledTimes(3);
+  });
+
+  it('should call the setState callback even if shouldComponentUpdate = false', done => {
+    const mockFn = jest.fn().mockReturnValue(false);
+
+    const Component = React.memo(
+      class Component extends React.Component {
+        constructor(props, context) {
+          super(props, context);
+          this.state = {
+            hasUpdatedState: false,
+          };
+        }
+
+        shouldComponentUpdate() {
+          return mockFn();
+        }
+
+        render() {
+          return <div>{this.state.hasUpdatedState}</div>;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<Component />);
+
+    const mountedInstance = shallowRenderer.getMountedInstance();
+    mountedInstance.setState({hasUpdatedState: true}, () => {
+      expect(mockFn).toBeCalled();
+      expect(mountedInstance.state.hasUpdatedState).toBe(true);
+      done();
+    });
+  });
+
+  it('throws usefully when rendering badly-typed elements', () => {
+    const shallowRenderer = createRenderer();
+
+    const renderAndVerifyWarningAndError = (Component, typeString) => {
+      expect(() => {
+        expect(() => shallowRenderer.render(<Component />)).toWarnDev(
+          'React.createElement: type is invalid -- expected a string ' +
+            '(for built-in components) or a class/function (for composite components) ' +
+            `but got: ${typeString}.`,
+        );
+      }).toThrowError(
+        'ReactShallowRenderer render(): Shallow rendering works only with custom ' +
+          `components, but the provided element type was \`${typeString}\`.`,
+      );
+    };
+
+    renderAndVerifyWarningAndError(undefined, 'undefined');
+    renderAndVerifyWarningAndError(null, 'null');
+    renderAndVerifyWarningAndError([], 'array');
+    renderAndVerifyWarningAndError({}, 'object');
+  });
+
+  it('should have initial state of null if not defined', () => {
+    const SomeComponent = React.memo(
+      class SomeComponent extends React.Component {
+        render() {
+          return <span />;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<SomeComponent />);
+
+    expect(shallowRenderer.getMountedInstance().state).toBeNull();
+  });
+
+  it('should invoke both deprecated and new lifecycles if both are present', () => {
+    const log = [];
+
+    const Component = React.memo(
+      class Component extends React.Component {
+        componentWillMount() {
+          log.push('componentWillMount');
+        }
+        componentWillReceiveProps() {
+          log.push('componentWillReceiveProps');
+        }
+        componentWillUpdate() {
+          log.push('componentWillUpdate');
+        }
+        UNSAFE_componentWillMount() {
+          log.push('UNSAFE_componentWillMount');
+        }
+        UNSAFE_componentWillReceiveProps() {
+          log.push('UNSAFE_componentWillReceiveProps');
+        }
+        UNSAFE_componentWillUpdate() {
+          log.push('UNSAFE_componentWillUpdate');
+        }
+        render() {
+          return null;
+        }
+      },
+    );
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<Component foo="bar" />);
+    expect(log).toEqual(['componentWillMount', 'UNSAFE_componentWillMount']);
+
+    log.length = 0;
+
+    shallowRenderer.render(<Component foo="baz" />);
+    expect(log).toEqual([
+      'componentWillReceiveProps',
+      'UNSAFE_componentWillReceiveProps',
+      'componentWillUpdate',
+      'UNSAFE_componentWillUpdate',
+    ]);
+  });
+
+  it('should stop the update when setState returns null or undefined', () => {
+    const log = [];
+    let instance;
+    const Component = React.memo(
+      class Component extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = {
+            count: 0,
+          };
+        }
+        render() {
+          log.push('render');
+          instance = this;
+          return null;
+        }
+      },
+    );
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<Component />);
+    log.length = 0;
+    instance.setState(() => null);
+    instance.setState(() => undefined);
+    instance.setState(null);
+    instance.setState(undefined);
+    expect(log).toEqual([]);
+    instance.setState(state => ({count: state.count + 1}));
+    expect(log).toEqual(['render']);
+  });
+
+  it('should not get this in a function component', () => {
+    const logs = [];
+    const Foo = React.memo(function Foo() {
+      logs.push(this);
+      return <div>foo</div>;
+    });
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<Foo foo="bar" />);
+    expect(logs).toEqual([undefined]);
+  });
+});