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

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

@@ -44,7 +44,7 @@ class ReactShallowRenderer {
       'ReactShallowRenderer render(): Invalid component element.%s',
       typeof element === 'function'
         ? ' Instead of passing a component class, make sure to instantiate ' +
-          'it by passing it to React.createElement.'
+            'it by passing it to React.createElement.'
         : '',
     );
     // Show a special message for host elements since it's a common case.
@@ -139,6 +139,7 @@ class ReactShallowRenderer {
     }
 
     this._rendered = this._instance.render();
+    this._updater._invokeCallback();
     // Intentionally do not call componentDidMount()
     // because DOM refs are not available.
   }
@@ -183,6 +184,7 @@ class ReactShallowRenderer {
 
     if (shouldUpdate) {
       this._rendered = this._instance.render();
+      this._updater._invokeCallback();
     }
     // Intentionally do not call componentDidUpdate()
     // because DOM refs are not available.
@@ -192,31 +194,41 @@ class ReactShallowRenderer {
 class Updater {
   constructor(renderer) {
     this._renderer = renderer;
+    this._callback = null;
+    this._publicInstance = null;
+  }
+
+  _updateCallback(callback, publicInstance) {
+    if (typeof callback === 'function') {
+      this._callback = callback;
+      this._publicInstance = publicInstance;
+    }
+  }
+
+  _invokeCallback() {
+    if (typeof this._callback === 'function' && this._publicInstance) {
+      this._callback.call(this._publicInstance);
+    }
   }
 
   isMounted(publicInstance) {
     return !!this._renderer._element;
   }
 
   enqueueForceUpdate(publicInstance, callback, callerName) {
+    this._updateCallback(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._updateCallback(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._updateCallback(callback, publicInstance);
     const currentState = this._renderer._newState || publicInstance.state;
 
     if (typeof partialState === 'function') {
@@ -229,10 +241,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,34 @@ describe('ReactShallowRenderer', () => {
     expect(result).toEqual(<div>baz:bar</div>);
   });
 
+  it('this.state should be updated on setState callback inside componentWillMount', function() {
+    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('throws usefully when rendering badly-typed elements', () => {
     spyOn(console, 'error');
     const shallowRenderer = createRenderer();