Add unstable_batchedUpdates into ShallowRenderer

koba04 · Aug 2, 2017 · ca78c92 · ca78c92
1 parent c5d1558
commit ca78c92
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 18 deletions.
diff --git a/src/renderers/testing/ReactShallowRendererEntry.js b/src/renderers/testing/ReactShallowRendererEntry.js
@@ -74,6 +74,10 @@ class ReactShallowRenderer {
     return this._rendered;
   }
 
+  unstable_batchedUpdates(fn) {
+    this._updater.batchedUpdates(fn);
+  }
+
   render(element, context = emptyObject) {
     invariant(
       React.isValidElement(element),
@@ -224,44 +228,73 @@ class ReactShallowRenderer {
 class Updater {
   constructor(renderer) {
     this._renderer = renderer;
+    this._isInBatched = false;
+    this._updateQueue = [];
+    this._callbackQueue = [];
+  }
+
+  batchedUpdates(fn) {
+    if (this._isInBatched) {
+      fn();
+      return;
+    }
+    this._isInBatched = true;
+    fn();
+    this._flushUpdates();
+    this._isInBatched = false;
   }
 
   isMounted(publicInstance) {
     return !!this._renderer._element;
   }
 
   enqueueForceUpdate(publicInstance, callback, callerName) {
-    this._renderer.render(this._renderer._element, this._renderer._context);
-
     if (typeof callback === 'function') {
-      callback.call(publicInstance);
+      this._callbackQueue.push(() => callback.call(publicInstance));
+    }
+
+    if (!this._isInBatched) {
+      this._flushUpdates();
     }
   }
 
   enqueueReplaceState(publicInstance, completeState, callback, callerName) {
-    this._renderer._newState = completeState;
-    this._renderer.render(this._renderer._element, this._renderer._context);
+    this._updateQueue.push(() => {
+      this._renderer._newState = completeState;
+    });
 
     if (typeof callback === 'function') {
-      callback.call(publicInstance);
+      this._callbackQueue.push(() => callback.call(publicInstance));
+    }
+    if (!this._isInBatched) {
+      this._flushUpdates();
     }
   }
 
   enqueueSetState(publicInstance, partialState, callback, callerName) {
-    if (typeof partialState === 'function') {
-      partialState = partialState(publicInstance.state, publicInstance.props);
+    this._updateQueue.push(() => {
+      if (typeof partialState === 'function') {
+        partialState = partialState(publicInstance.state, publicInstance.props);
+      }
+      this._renderer._newState = {
+        ...publicInstance.state,
+        ...partialState,
+      };
+    });
+    if (typeof callback === 'function') {
+      this._callbackQueue.push(() => callback.call(publicInstance));
     }
+    if (!this._isInBatched) {
+      this._flushUpdates();
+    }
+  }
 
-    this._renderer._newState = {
-      ...publicInstance.state,
-      ...partialState,
-    };
-
+  _flushUpdates() {
+    this._updateQueue.forEach(queue => queue());
     this._renderer.render(this._renderer._element, this._renderer._context);
-
-    if (typeof callback === 'function') {
-      callback.call(publicInstance);
-    }
+    this._callbackQueue.forEach(callback => callback());
+    this._updateQueue = [];
+    this._callbackQueue = [];
   }
 }
 

diff --git a/src/renderers/testing/__tests__/ReactShallowRenderer-test.js b/src/renderers/testing/__tests__/ReactShallowRenderer-test.js
@@ -15,7 +15,7 @@ let createRenderer;
 let PropTypes;
 let React;
 
-describe('ReactTestUtils', () => {
+describe('ReactShallowRenderer', () => {
   beforeEach(() => {
     createRenderer = require('react-test-renderer/shallow').createRenderer;
     PropTypes = require('prop-types');
@@ -658,4 +658,26 @@ describe('ReactTestUtils', () => {
     result = shallowRenderer.render(cloned);
     expect(result).toEqual(<div>baz:bar</div>);
   });
+
+  it('should enable batchedUpdates', () => {
+    let renderCount = 0;
+    class App extends React.Component {
+      state = {count: 0};
+      render() {
+        ++renderCount;
+        return <div>hoge</div>;
+      }
+    }
+
+    const shallowRenderer = createRenderer();
+    shallowRenderer.render(<App />);
+    const instance = shallowRenderer.getMountedInstance();
+    renderCount = 0;
+    shallowRenderer.unstable_batchedUpdates(() => {
+      instance.setState({count: instance.state.count + 1});
+      instance.setState({count: instance.state.count + 1});
+    });
+    expect(renderCount).toEqual(1);
+    expect(shallowRenderer.getMountedInstance().state.count).toEqual(1);
+  });
 });