Fix to reset the forceUpdate flag after the update (#11440)

* Fix to reset the forceUpdate flag after the update

* Add support for PureComponent and mirror real behavior closer

packages/react-test-renderer/src/ReactShallowRenderer.js

@@ -11,6 +11,7 @@ import describeComponentFrame from 'shared/describeComponentFrame';
 import getComponentName from 'shared/getComponentName';
 import emptyObject from 'fbjs/lib/emptyObject';
 import invariant from 'fbjs/lib/invariant';
+import shallowEqual from 'fbjs/lib/shallowEqual';
 import checkPropTypes from 'prop-types/checkPropTypes';
 
 class ReactShallowRenderer {
@@ -63,7 +64,7 @@ class ReactShallowRenderer {
     this._context = context;
 
     if (this._instance) {
-      this._updateClassComponent(element.props, context);
+      this._updateClassComponent(element.type, element.props, context);
     } else {
       if (shouldConstruct(element.type)) {
         this._instance = new element.type(
@@ -133,7 +134,8 @@ class ReactShallowRenderer {
     // because DOM refs are not available.
   }
 
-  _updateClassComponent(props, context) {
+  _updateClassComponent(type, props, context) {
+    const oldState = this._instance.state || emptyObject;
     const oldProps = this._instance.props;
 
     if (
@@ -145,33 +147,40 @@ class ReactShallowRenderer {
 
     // Read state after cWRP in case it calls setState
     // Fallback to previous instance state to support rendering React.cloneElement()
+    // TODO: the cloneElement() use case is broken and should be removed
+    // https://github.com/facebook/react/issues/11441
     const state = this._newState || this._instance.state || emptyObject;
 
-    if (
-      typeof this._instance.shouldComponentUpdate === 'function' &&
-      !this._forcedUpdate
-    ) {
-      if (
-        this._instance.shouldComponentUpdate(props, state, context) === false
-      ) {
-        this._instance.context = context;
-        this._instance.props = props;
-        this._instance.state = state;
-        this._forcedUpdate = false;
-
-        return;
-      }
+    let shouldUpdate = true;
+    if (this._forcedUpdate) {
+      shouldUpdate = true;
+      this._forcedUpdate = false;
+    } else if (typeof this._instance.shouldComponentUpdate === 'function') {
+      shouldUpdate = !!this._instance.shouldComponentUpdate(
+        props,
+        state,
+        context,
+      );
+    } else if (type && type.prototype && type.prototype.isPureReactComponent) {
+      // TODO: we can remove the type existence check when we fix this:
+      // https://github.com/facebook/react/issues/11441
+      shouldUpdate =
+        !shallowEqual(oldProps, props) || !shallowEqual(oldState, state);
     }
 
-    if (typeof this._instance.componentWillUpdate === 'function') {
-      this._instance.componentWillUpdate(props, state, context);
+    if (shouldUpdate) {
+      if (typeof this._instance.componentWillUpdate === 'function') {
+        this._instance.componentWillUpdate(props, state, context);
+      }
     }
 
     this._instance.context = context;
     this._instance.props = props;
     this._instance.state = state;
 
-    this._rendered = this._instance.render();
+    if (shouldUpdate) {
+      this._rendered = this._instance.render();
+    }
     // Intentionally do not call componentDidUpdate()
     // because DOM refs are not available.
   }

packages/react-test-renderer/src/__tests__/ReactShallowRenderer-test.js

@@ -122,24 +122,64 @@ describe('ReactShallowRenderer', () => {
     expect(shallowRenderer.getRenderOutput()).toEqual(<div>2</div>);
   });
 
+  it('should enable PureComponent to prevent a re-render', () => {
+    let renderCounter = 0;
+    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;
     class SimpleComponent extends React.Component {
+      state = {count: 1};
       shouldComponentUpdate() {
         scuCounter++;
+        return false;
       }
       render() {
-        return <div />;
+        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', () => {