Bug fix - SetState callback called before component state is updated in ReactShallowRenderer

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

@@ -338,4 +338,42 @@ describe('ReactDOM', () => {
     ];
     expect(actual).toEqual(expected);
   });
+
+  it('should call the setState callback even if shouldComponentUpdate = false', () => {
+    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);
+    });
+  });
 });

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

@@ -103,6 +103,7 @@ class ReactShallowRenderer {
     }
 
     this._rendering = false;
+    this._updater._invokeCallbackIfNecessary();
 
     return this.getRenderOutput();
   }
@@ -192,31 +193,44 @@ class ReactShallowRenderer {
 class Updater {
   constructor(renderer) {
     this._renderer = renderer;
+    this._callback = null;
+    this._publicInstance = null;
+  }
+
+  _enqueueCallback(callback, publicInstance) {
+    if (typeof callback === 'function') {
+      this._callback = callback;
+      this._publicInstance = publicInstance;
+    }
+  }
+
+  _invokeCallbackIfNecessary() {
+    if (typeof this._callback === 'function' && this._publicInstance) {
+      const {_callback, _publicInstance} = this;
+      this._callback = null;
+      this._publicInstance = null;
+      _callback.call(_publicInstance);
+    }
   }
 
   isMounted(publicInstance) {
     return !!this._renderer._element;
   }
 
   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 +243,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

@@ -845,6 +845,64 @@ 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 call the setState callback even if shouldComponentUpdate = false', () => {
+    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);
+    });
+  });
+
   it('throws usefully when rendering badly-typed elements', () => {
     spyOn(console, 'error');
     const shallowRenderer = createRenderer();