From 9c21e2f3c447776bb38a71b12b60a1c3218edf97 Mon Sep 17 00:00:00 2001
From: Scott Feeney <smf@fb.com>
Date: Sat, 28 Nov 2015 16:42:36 -0800
Subject: [PATCH] shallow render: fix setState in componentWillMount

---
 src/test/ReactTestUtils.js                | 10 +++++++---
 src/test/__tests__/ReactTestUtils-test.js | 15 +++++++++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/test/ReactTestUtils.js b/src/test/ReactTestUtils.js
index 467b6682c2c8..8b5eb5fecdce 100644
--- a/src/test/ReactTestUtils.js
+++ b/src/test/ReactTestUtils.js
@@ -427,13 +427,17 @@ ReactShallowRenderer.prototype.render = function(element, context) {
   if (!context) {
     context = emptyObject;
   }
-  var transaction = ReactUpdates.ReactReconcileTransaction.getPooled(true);
-  this._render(element, transaction, context);
-  ReactUpdates.ReactReconcileTransaction.release(transaction);
+  ReactUpdates.batchedUpdates(_batchedRender, this, element, context);
 
   return this.getRenderOutput();
 };
 
+function _batchedRender(renderer, element, context) {
+  var transaction = ReactUpdates.ReactReconcileTransaction.getPooled(true);
+  renderer._render(element, transaction, context);
+  ReactUpdates.ReactReconcileTransaction.release(transaction);
+}
+
 ReactShallowRenderer.prototype.getRenderOutput = function() {
   return (
     (this._instance && this._instance._renderedComponent &&
diff --git a/src/test/__tests__/ReactTestUtils-test.js b/src/test/__tests__/ReactTestUtils-test.js
index 86fc3d13f920..1fb6956ffa39 100644
--- a/src/test/__tests__/ReactTestUtils-test.js
+++ b/src/test/__tests__/ReactTestUtils-test.js
@@ -217,6 +217,21 @@ describe('ReactTestUtils', function() {
     expect(result.props.className).toEqual('clicked');
   });
 
+  it('can setState in componentWillMount when shallow rendering', function() {
+    var SimpleComponent = React.createClass({
+      componentWillMount() {
+        this.setState({groovy: 'doovy'});
+      },
+      render() {
+        return <div>{this.state.groovy}</div>;
+      },
+    });
+
+    var shallowRenderer = ReactTestUtils.createRenderer();
+    var result = shallowRenderer.render(<SimpleComponent />);
+    expect(result).toEqual(<div>doovy</div>);
+  });
+
   it('can pass context when shallowly rendering', function() {
     var SimpleComponent = React.createClass({
       contextTypes: {