From 066aefdaed547bb253f211ac9c7dda47a2c15458 Mon Sep 17 00:00:00 2001
From: Toru Kobayashi <koba0004@gmail.com>
Date: Fri, 22 Apr 2016 18:42:46 +0900
Subject: [PATCH] ShallowRenderer supports batched updates

---
 src/ShallowWrapper.js       |  5 ++++-
 src/react-compat.js         |  2 ++
 test/ShallowWrapper-spec.js | 28 ++++++++++++++++++++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/ShallowWrapper.js b/src/ShallowWrapper.js
index 4589fcd79..dc8e19cf6 100644
--- a/src/ShallowWrapper.js
+++ b/src/ShallowWrapper.js
@@ -25,6 +25,7 @@ import {
 import {
   createShallowRenderer,
   renderToStaticMarkup,
+  batchedUpdates,
 } from './react-compat';
 
 /**
@@ -352,7 +353,9 @@ export default class ShallowWrapper {
       withSetStateAllowed(() => {
         // TODO(lmr): create/use synthetic events
         // TODO(lmr): emulate React's event propagation
-        handler(...args);
+        batchedUpdates(() => {
+          handler(...args);
+        });
         this.root.update();
       });
     }
diff --git a/src/react-compat.js b/src/react-compat.js
index 059aa3e7f..f434853e3 100644
--- a/src/react-compat.js
+++ b/src/react-compat.js
@@ -12,6 +12,7 @@ let renderWithOptions;
 let unmountComponentAtNode;
 
 const React = require('react');
+const batchedUpdates = require('react/lib/ReactUpdates').batchedUpdates;
 
 if (REACT013) {
   renderToStaticMarkup = React.renderToStaticMarkup;
@@ -163,4 +164,5 @@ export {
   childrenToArray,
   renderWithOptions,
   unmountComponentAtNode,
+  batchedUpdates,
 };
diff --git a/test/ShallowWrapper-spec.js b/test/ShallowWrapper-spec.js
index 475bfb93d..8ea437ead 100644
--- a/test/ShallowWrapper-spec.js
+++ b/test/ShallowWrapper-spec.js
@@ -974,6 +974,34 @@ describe('shallow', () => {
       });
     });
 
+    it('should be batched updates', () => {
+      let renderCount = 0;
+      class Foo extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = {
+            count: 0,
+          };
+          this.onClick = this.onClick.bind(this);
+        }
+        onClick() {
+          this.setState({ count: this.state.count + 1 });
+          this.setState({ count: this.state.count + 1 });
+        }
+        render() {
+          ++renderCount;
+          return (
+            <a onClick={this.onClick}>{this.state.count}</a>
+          );
+        }
+      }
+
+      const wrapper = shallow(<Foo />);
+      wrapper.simulate('click');
+      expect(wrapper.text()).to.equal('1');
+      expect(renderCount).to.equal(2);
+    });
+
   });
 
   describe('.setState(newState)', () => {