Fix shallow renderer callbacks (#10106)

* Add failing test to show that shallow test renderer doesn't call setState's callback arg
* Record tests
* Fix shallow renderer's setState/replaceState/forceUpdate to execute any callbacks passed. (#10089)
* Ensure shallow renderer callbacks are called with the correct  binding.

scripts/fiber/tests-passing.txt

@@ -1955,6 +1955,9 @@ src/renderers/dom/test/__tests__/ReactTestUtils-test.js
 * can setState in componentWillMount when shallow rendering
 * can setState in componentWillReceiveProps when shallow rendering
 * can setState with an updater function
+* can setState with a callback
+* can replaceState with a callback
+* can forceUpdate with a callback
 * can pass context when shallowly rendering
 * should track context across updates
 * can fail context when shallowly rendering

src/renderers/dom/test/__tests__/ReactTestUtils-test.js

@@ -405,6 +405,108 @@ describe('ReactTestUtils', () => {
     expect(result.props.children).toEqual(2);
   });
 
+  it('can setState with a callback', () => {
+    let instance;
+
+    class SimpleComponent extends React.Component {
+      state = {
+        counter: 0,
+      };
+      render() {
+        instance = this;
+        return (
+          <p>
+            {this.state.counter}
+          </p>
+        );
+      }
+    }
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SimpleComponent />);
+    expect(result.props.children).toBe(0);
+
+    const callback = jest.fn(function() {
+      expect(this).toBe(instance);
+    });
+
+    instance.setState({counter: 1}, callback);
+
+    const updated = shallowRenderer.getRenderOutput();
+    expect(updated.props.children).toBe(1);
+    expect(callback).toHaveBeenCalled();
+  });
+
+  it('can replaceState with a callback', () => {
+    let instance;
+
+    class SimpleComponent extends React.Component {
+      state = {
+        counter: 0,
+      };
+      render() {
+        instance = this;
+        return (
+          <p>
+            {this.state.counter}
+          </p>
+        );
+      }
+    }
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SimpleComponent />);
+    expect(result.props.children).toBe(0);
+
+    const callback = jest.fn(function() {
+      expect(this).toBe(instance);
+    });
+
+    // No longer a public API, but we can test that it works internally by
+    // reaching into the updater.
+    shallowRenderer._updater.enqueueReplaceState(
+      instance,
+      {counter: 1},
+      callback,
+    );
+
+    const updated = shallowRenderer.getRenderOutput();
+    expect(updated.props.children).toBe(1);
+    expect(callback).toHaveBeenCalled();
+  });
+
+  it('can forceUpdate with a callback', () => {
+    let instance;
+
+    class SimpleComponent extends React.Component {
+      state = {
+        counter: 0,
+      };
+      render() {
+        instance = this;
+        return (
+          <p>
+            {this.state.counter}
+          </p>
+        );
+      }
+    }
+
+    const shallowRenderer = createRenderer();
+    const result = shallowRenderer.render(<SimpleComponent />);
+    expect(result.props.children).toBe(0);
+
+    const callback = jest.fn(function() {
+      expect(this).toBe(instance);
+    });
+
+    instance.forceUpdate(callback);
+
+    const updated = shallowRenderer.getRenderOutput();
+    expect(updated.props.children).toBe(0);
+    expect(callback).toHaveBeenCalled();
+  });
+
   it('can pass context when shallowly rendering', () => {
     class SimpleComponent extends React.Component {
       static contextTypes = {

src/renderers/testing/ReactShallowRendererEntry.js

@@ -201,11 +201,19 @@ class Updater {
 
   enqueueForceUpdate(publicInstance, callback, callerName) {
     this._renderer.render(this._renderer._element, this._renderer._context);
+
+    if (typeof callback === 'function') {
+      callback.call(publicInstance);
+    }
   }
 
   enqueueReplaceState(publicInstance, completeState, callback, callerName) {
     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) {
@@ -219,6 +227,10 @@ class Updater {
     };
 
     this._renderer.render(this._renderer._element, this._renderer._context);
+
+    if (typeof callback === 'function') {
+      callback.call(publicInstance);
+    }
   }
 }