**Do you want to request a *feature* or report a *bug*?**
report a bug

**What is the current behavior?**
When I call setState with a callback in a test using react shallow renderer (via enzyme), the callback gets called and `this.state` is still the old state.

EDIT: This seems limited to componentWillMount

**If the current behavior is a bug, please provide the steps to reproduce and if possible a minimal demo of the problem via https://jsfiddle.net or similar (template for React 16: https://jsfiddle.net/Luktwrdm/, template for React 15: https://jsfiddle.net/hmbg7e9w/).**

EDIT: Reproduce with this https://github.com/bdwain/setstate-callback-bug

**What is the expected behavior?**
When the setState callback gets called, it should have access to the new state.

**Which versions of React, and which browser / OS are affected by this issue? Did this work in previous versions of React?**
Version 16.
It worked in v15 with enzyme 2.

EDIT: I think this is because in componentWillMount, [this line](https://github.com/facebook/react/blob/master/packages/react-test-renderer/src/ReactShallowRenderer.js#L67) causes the render method in the shallow wrapper to return early, before it updates the state.
it('should warn about `setState` on unmounted components', () => {
    spyOnDev(console, 'error');
    var container = document.createElement('div');
    document.body.appendChild(container);
    var renders = 0;
    class Component extends React.Component {
      state = {value: 0};
      render() {
        renders++;
        return <div />;
      }
    }
    var instance = <Component />;
    expect(instance.setState).not.toBeDefined();
    instance = ReactDOM.render(instance, container);
    expect(renders).toBe(1);
    instance.setState({value: 1});
    if (__DEV__) {
      expect(console.error.calls.count()).toBe(0);
    }
    expect(renders).toBe(2);
    ReactDOM.unmountComponentAtNode(container);
    instance.setState({value: 2});
    expect(renders).toBe(2);
    if (__DEV__) {
      expect(console.error.calls.count()).toBe(1);
      expect(console.error.calls.argsFor(0)[0]).toContain(
        'Can only update a mounted or mounting component. This usually means ' +
          'you called setState, replaceState, or forceUpdate on an unmounted ' +
          'component. This is a no-op.\n\nPlease check the code for the ' +