**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.
var ReactDOM;
var ReactDOMServer;
var ReactCurrentOwner;
var ReactTestUtils;
var PropTypes;
var shallowEqual;
var shallowCompare;
describe('ReactCompositeComponent', () => {
  beforeEach(() => {
    jest.resetModules();
    React = require('react');
    ReactDOM = require('react-dom');
    ReactDOMServer = require('react-dom/server');
    ReactCurrentOwner = require('react')
      .__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.ReactCurrentOwner;
    ReactTestUtils = require('react-dom/test-utils');
    PropTypes = require('prop-types');
    shallowEqual = require('fbjs/lib/shallowEqual');
    shallowCompare = function(instance, nextProps, nextState) {
      return (
        !shallowEqual(instance.props, nextProps) ||
        !shallowEqual(instance.state, nextState)
      );
    };
    MorphingComponent = class extends React.Component {
      state = {activated: false};
      _toggleActivatedState = () => {
        this.setState({activated: !this.state.activated});
      };
      render() {