**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.
expect(moo.state.amIImmutable()).toBe(true);
    var secondState = new NotActuallyImmutable('second');
    moo._replaceState(secondState);
    expect(moo.state.str).toBe('second');
    expect(moo.state.amIImmutable()).toBe(true);
    expect(moo.state).toBe(secondState);
    moo.setState({str: 'third'});
    expect(moo.state.str).toBe('third');
    
    expect(moo.state.amIImmutable).toBe(undefined);
    
    var fifthState = new NotActuallyImmutable('fifth');
    ReactDOM.unstable_batchedUpdates(function() {
      moo.setState({str: 'fourth'});
      moo._replaceState(fifthState);
    });
    expect(moo.state).toBe(fifthState);
    
    var sixthState = new NotActuallyImmutable('sixth');
    ReactDOM.unstable_batchedUpdates(function() {
      moo._replaceState(sixthState);
      moo.setState({str: 'seventh'});
    });
    expect(moo.state.str).toBe('seventh');
    expect(moo.state.amIImmutable).toBe(undefined);
  });
  it('should not warn about unmounting during unmounting', () => {
    var container = document.createElement('div');
    var layer = document.createElement('div');
    class Component extends React.Component {