Bug fix - SetState callback called before component state is updated in ReactShallowRenderer (#11507)

* Create test to verify ReactShallowRenderer bug (#11496)

* Fix ReactShallowRenderer callback bug on componentWillMount (#11496)

* Improve fnction naming and clean up queued callback before call

* Run prettier on ReactShallowRenderer.js

* Consolidate callback call on ReactShallowRenderer.js

* Ensure callback behavior is similar between ReactDOM and ReactShallowRenderer

* Fix Code Review requests (#11507)

* Move test to ReactCompositeComponent

* Verify the callback gets called

* Ensure multiple callbacks are correctly handled on ReactShallowRenderer

* Ensure the setState callback is called inside componentWillMount (ReactDOM)

* Clear ReactShallowRenderer callback queue before actually calling the callbacks

* Add test for multiple callbacks on ReactShallowRenderer

* Ensure the ReactShallowRenderer callback queue is cleared after invoking callbacks

* Remove references to internal fields on ReactShallowRenderer test

packages/react-dom/src/__tests__/ReactCompositeComponent-test.js

@@ -1634,6 +1634,73 @@ describe('ReactCompositeComponent', () => {
     expect(mockArgs.length).toEqual(0);
   });
 
+  it('this.state should be updated on setState callback inside componentWillMount', () => {
+    const div = document.createElement('div');
+    let stateSuccessfullyUpdated = false;
+
+    class Component extends React.Component {
+      constructor(props, context) {
+        super(props, context);
+        this.state = {
+          hasUpdatedState: false,
+        };
+      }
+
+      componentWillMount() {
+        this.setState(
+          {hasUpdatedState: true},
+          () => (stateSuccessfullyUpdated = this.state.hasUpdatedState),
+        );
+      }
+
+      render() {
+        return <div>{this.props.children}</div>;
+      }
+    }
+
+    ReactDOM.render(<Component />, div);
+    expect(stateSuccessfullyUpdated).toBe(true);
+  });
+
+  it('should call the setState callback even if shouldComponentUpdate = false', done => {
+    const mockFn = jest.fn().mockReturnValue(false);
+    const div = document.createElement('div');
+
+    let instance;
+
+    class Component extends React.Component {
+      constructor(props, context) {
+        super(props, context);
+        this.state = {
+          hasUpdatedState: false,
+        };
+      }
+
+      componentWillMount() {
+        instance = this;
+      }
+
+      shouldComponentUpdate() {
+        return mockFn();
+      }
+
+      render() {
+        return <div>{this.state.hasUpdatedState}</div>;
+      }
+    }
+
+    ReactDOM.render(<Component />, div);
+
+    expect(instance).toBeDefined();
+    expect(mockFn).not.toBeCalled();
+
+    instance.setState({hasUpdatedState: true}, () => {
+      expect(mockFn).toBeCalled();
+      expect(instance.state.hasUpdatedState).toBe(true);
+      done();
+    });
+  });
+
   it('should return a meaningful warning when constructor is returned', () => {
     spyOnDev(console, 'error');
     class RenderTextInvalidConstructor extends React.Component {

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

@@ -103,6 +103,7 @@ class ReactShallowRenderer {
     }
 
     this._rendering = false;
+    this._updater._invokeCallbacks();
 
     return this.getRenderOutput();
   }
@@ -192,6 +193,25 @@ class ReactShallowRenderer {
 class Updater {
   constructor(renderer) {
     this._renderer = renderer;
+    this._callbacks = [];
+  }
+
+  _enqueueCallback(callback, publicInstance) {
+    if (typeof callback === 'function' && publicInstance) {
+      this._callbacks.push({
+        callback,
+        publicInstance,
+      });
+    }
+  }
+
+  _invokeCallbacks() {
+    const callbacks = this._callbacks;
+    this._callbacks = [];
+
+    callbacks.forEach(({callback, publicInstance}) => {
+      callback.call(publicInstance);
+    });
   }
 
   isMounted(publicInstance) {
@@ -199,24 +219,19 @@ class Updater {
   }
 
   enqueueForceUpdate(publicInstance, callback, callerName) {
+    this._enqueueCallback(callback, publicInstance);
     this._renderer._forcedUpdate = true;
     this._renderer.render(this._renderer._element, this._renderer._context);
-
-    if (typeof callback === 'function') {
-      callback.call(publicInstance);
-    }
   }
 
   enqueueReplaceState(publicInstance, completeState, callback, callerName) {
+    this._enqueueCallback(callback, publicInstance);
     this._renderer._newState = completeState;
     this._renderer.render(this._renderer._element, this._renderer._context);
-
-    if (typeof callback === 'function') {
-      callback.call(publicInstance);
-    }
   }
 
   enqueueSetState(publicInstance, partialState, callback, callerName) {
+    this._enqueueCallback(callback, publicInstance);
     const currentState = this._renderer._newState || publicInstance.state;
 
     if (typeof partialState === 'function') {
@@ -229,10 +244,6 @@ class Updater {
     };
 
     this._renderer.render(this._renderer._element, this._renderer._context);
-
-    if (typeof callback === 'function') {
-      callback.call(publicInstance);
-    }
   }
 }
 

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

@@ -849,6 +849,97 @@ describe('ReactShallowRenderer', () => {
     expect(result).toEqual(<div>baz:bar</div>);
   });
 
+  it('this.state should be updated on setState callback inside componentWillMount', () => {
+    let stateSuccessfullyUpdated = false;
+
+    class Component extends React.Component {
+      constructor(props, context) {
+        super(props, context);
+        this.state = {
+          hasUpdatedState: false,
+        };
+      }
+
+      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();
+
+    class Component extends React.Component {
+      constructor(props, context) {
+        super(props, context);
+        this.state = {
+          foo: 'foo',
+        };
+      }
+
+      componentWillMount() {
+        this.setState({foo: 'bar'}, () => mockFn());
+        this.setState({foo: 'foobar'}, () => mockFn());
+      }
+
+      render() {
+        return <div>{this.state.foo}</div>;
+      }
+    }
+
+    shallowRenderer.render(<Component />);
+
+    expect(mockFn.mock.calls.length).toBe(2);
+
+    // Ensure the callback queue is cleared after the callbacks are invoked
+    const mountedInstance = shallowRenderer.getMountedInstance();
+    mountedInstance.setState({foo: 'bar'}, () => mockFn());
+    expect(mockFn.mock.calls.length).toBe(3);
+  });
+
+  it('should call the setState callback even if shouldComponentUpdate = false', done => {
+    const mockFn = jest.fn().mockReturnValue(false);
+
+    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', () => {
     spyOnDev(console, 'error');
     const shallowRenderer = createRenderer();