From f3510cfc463f6ed94654b414c5ae305219eb3d5a Mon Sep 17 00:00:00 2001
From: Alex Cordeiro <accordeiro@users.noreply.github.com>
Date: Wed, 22 Nov 2017 20:15:11 -0200
Subject: [PATCH] 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
---
 .../__tests__/ReactCompositeComponent-test.js | 67 ++++++++++++++
 .../src/ReactShallowRenderer.js               | 35 ++++---
 .../__tests__/ReactShallowRenderer-test.js    | 91 +++++++++++++++++++
 3 files changed, 181 insertions(+), 12 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactCompositeComponent-test.js b/packages/react-dom/src/__tests__/ReactCompositeComponent-test.js
index 5797e4b9429340..1bac93ec046c10 100644
--- a/packages/react-dom/src/__tests__/ReactCompositeComponent-test.js
+++ b/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 {
diff --git a/packages/react-test-renderer/src/ReactShallowRenderer.js b/packages/react-test-renderer/src/ReactShallowRenderer.js
index 1958f43d7bd798..5cf10e7bb935e4 100644
--- a/packages/react-test-renderer/src/ReactShallowRenderer.js
+++ b/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);
-    }
   }
 }
 
diff --git a/packages/react-test-renderer/src/__tests__/ReactShallowRenderer-test.js b/packages/react-test-renderer/src/__tests__/ReactShallowRenderer-test.js
index 9faf6a0ff8e927..6d22d868ec0f5e 100644
--- a/packages/react-test-renderer/src/__tests__/ReactShallowRenderer-test.js
+++ b/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();