Fixes #509: pass callback on setState and setProps

docs/README.md

@@ -57,7 +57,7 @@
     * [render()](/docs/api/ShallowWrapper/render.md)
     * [setContext(context)](/docs/api/ShallowWrapper/setContext.md)
     * [setProps(nextProps)](/docs/api/ShallowWrapper/setProps.md)
-    * [setState(nextState)](/docs/api/ShallowWrapper/setState.md)
+    * [setState(nextState[, callback])](/docs/api/ShallowWrapper/setState.md)
     * [shallow([options])](/docs/api/ShallowWrapper/shallow.md)
     * [simulate(event[, data])](/docs/api/ShallowWrapper/simulate.md)
     * [some(selector)](/docs/api/ShallowWrapper/some.md)
@@ -110,8 +110,8 @@
     * [ref(refName)](/docs/api/ReactWrapper/ref.md)
     * [render()](/docs/api/ReactWrapper/render.md)
     * [setContext(context)](/docs/api/ReactWrapper/setContext.md)
-    * [setProps(nextProps)](/docs/api/ReactWrapper/setProps.md)
-    * [setState(nextState)](/docs/api/ReactWrapper/setState.md)
+    * [setProps(nextProps[, callback])](/docs/api/ReactWrapper/setProps.md)
+    * [setState(nextState[, callback])](/docs/api/ReactWrapper/setState.md)
     * [simulate(event[, data])](/docs/api/ReactWrapper/simulate.md)
     * [some(selector)](/docs/api/ReactWrapper/some.md)
     * [someWhere(predicate)](/docs/api/ReactWrapper/someWhere.md)

docs/api/ReactWrapper/setProps.md

@@ -1,4 +1,4 @@
-# `.setProps(props) => Self`
+# `.setProps(props[, callback]) => Self`
 
 A method that sets the props of the root component, and re-renders. Useful for when you are
 wanting to test how the component behaves over time with changing props. Calling this, for
@@ -13,7 +13,7 @@ NOTE: can only be called on a wrapper instance that is also the root instance.
 #### Arguments
 
 1. `props` (`Object`): An object containing new props to merge in with the current state
-
+2. `callback` (`Function` [optional]): If provided, the callback function will be executed once setProps has completed
 
 
 #### Returns

docs/api/ReactWrapper/setState.md

@@ -1,4 +1,4 @@
-# `.setState(state) => Self`
+# `.setState(state[, callback]) => Self`
 
 A method to invoke `setState()` on the root component instance similar to how you might in the
 definition of the component, and re-renders.  This method is useful for testing your component
@@ -12,7 +12,7 @@ NOTE: can only be called on a wrapper instance that is also the root instance.
 #### Arguments
 
 1. `state` (`Object`): An object containing new state to merge in with the current state
-
+2. `callback` (`Function` [optional]): If provided, the callback function will be executed once setState has completed
 
 
 #### Returns

docs/api/ShallowWrapper/setState.md

@@ -1,4 +1,4 @@
-# `.setState(state) => Self`
+# `.setState(state[, callback]) => Self`
 
 A method to invoke `setState()` on the root component instance similar to how you might in the
 definition of the component, and re-renders.  This method is useful for testing your component
@@ -12,7 +12,7 @@ NOTE: can only be called on a wrapper instance that is also the root instance.
 #### Arguments
 
 1. `state` (`Object`): An object containing new state to merge in with the current state
-
+2. `callback` (`Function` [optional]): If provided, the callback function will be executed once setState has completed
 
 
 #### Returns

src/ReactWrapper.jsx

@@ -205,13 +205,14 @@ class ReactWrapper {
    * NOTE: can only be called on a wrapper instance that is also the root instance.
    *
    * @param {Object} props object
+   * @param {Function} cb - callback function
    * @returns {ReactWrapper}
    */
-  setProps(props) {
+  setProps(props, callback = undefined) {
     if (this.root !== this) {
       throw new Error('ReactWrapper::setProps() can only be called on the root');
     }
-    this.component.setChildProps(props);
+    this.component.setChildProps(props, callback);
     return this;
   }
 
@@ -225,13 +226,14 @@ class ReactWrapper {
    * NOTE: can only be called on a wrapper instance that is also the root instance.
    *
    * @param {Object} state to merge
+   * @param {Function} cb - callback function
    * @returns {ReactWrapper}
    */
-  setState(state) {
+  setState(state, callback = undefined) {
     if (this.root !== this) {
       throw new Error('ReactWrapper::setState() can only be called on the root');
     }
-    this.instance().setState(state);
+    this.instance().setState(state, callback);
     return this;
   }
 

src/ReactWrapperComponent.jsx

@@ -34,9 +34,9 @@ export default function createWrapperComponent(node, options = {}) {
       };
     },
 
-    setChildProps(newProps) {
+    setChildProps(newProps, callback = undefined) {
       const props = objectAssign({}, this.state.props, newProps);
-      this.setState({ props });
+      this.setState({ props }, callback);
     },
 
     setChildContext(context) {

src/ShallowWrapper.js

@@ -245,9 +245,10 @@ class ShallowWrapper {
    * NOTE: can only be called on a wrapper instance that is also the root instance.
    *
    * @param {Object} state to merge
+   * @param {Function} cb - callback function
    * @returns {ShallowWrapper}
    */
-  setState(state) {
+  setState(state, callback = undefined) {
     if (this.root !== this) {
       throw new Error('ShallowWrapper::setState() can only be called on the root');
     }
@@ -256,7 +257,7 @@ class ShallowWrapper {
     }
     this.single('setState', () => {
       withSetStateAllowed(() => {
-        this.instance().setState(state);
+        this.instance().setState(state, callback);
         this.update();
       });
     });

test/ReactWrapper-spec.jsx

@@ -3,6 +3,7 @@
 import React from 'react';
 import { expect } from 'chai';
 import sinon from 'sinon';
+import { batchedUpdates } from '../src/react-compat';
 
 import {
   describeWithDOM,
@@ -857,7 +858,7 @@ describeWithDOM('mount', () => {
 
   });
 
-  describe('.setProps(newProps)', () => {
+  describe('.setProps(newProps[, callback])', () => {
     it('should set props for a component multiple times', () => {
       class Foo extends React.Component {
         render() {
@@ -960,6 +961,27 @@ describeWithDOM('mount', () => {
       expect(setInvalidProps).to.throw(TypeError, similarException.message);
     });
 
+    it('should call the callback when setProps has completed', () => {
+      class Foo extends React.Component {
+        render() {
+          return (
+            <div className={this.props.id}>
+              {this.props.id}
+            </div>
+          );
+        }
+      }
+      const wrapper = mount(<Foo id="foo" />);
+      expect(wrapper.find('.foo').length).to.equal(1);
+
+      batchedUpdates(() => {
+        wrapper.setProps({ id: 'bar', foo: 'bla' }, () => {
+          expect(wrapper.find('.bar').length).to.equal(1);
+        });
+        expect(wrapper.find('.bar').length).to.equal(0);
+      });
+    });
+
     describeIf(!REACT013, 'stateless function components', () => {
       it('should set props for a component multiple times', () => {
         const Foo = props => (
@@ -1268,7 +1290,7 @@ describeWithDOM('mount', () => {
     });
   });
 
-  describe('.setState(newState)', () => {
+  describe('.setState(newState[, callback])', () => {
     it('should set the state of the root node', () => {
       class Foo extends React.Component {
         constructor(props) {
@@ -1305,6 +1327,63 @@ describeWithDOM('mount', () => {
       const wrapper = mount(<MySharona />);
       expect(wrapper.find('div').text()).to.equal('a');
     });
+
+    it('should call the callback when setState has completed', () => {
+      class Foo extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = { id: 'foo' };
+        }
+        render() {
+          return (
+            <div className={this.state.id} />
+          );
+        }
+      }
+      const wrapper = mount(<Foo />);
+      expect(wrapper.state()).to.eql({ id: 'foo' });
+      wrapper.setState({ id: 'bar' }, () => {
+        expect(wrapper.state()).to.eql({ id: 'bar' });
+      });
+    });
+
+    it('should throw error when cb is not a function', () => {
+      class Foo extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = { id: 'foo' };
+        }
+        render() {
+          return (
+            <div className={this.state.id} />
+          );
+        }
+      }
+      const wrapper = mount(<Foo />);
+      expect(wrapper.state()).to.eql({ id: 'foo' });
+      expect(() => wrapper.setState({ id: 'bar' }, 1)).to.throw(Error);
+    });
+
+    itIf(REACT15, 'should throw error when cb is not a function', () => {
+      class Foo extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = { id: 'foo' };
+        }
+        render() {
+          return (
+            <div className={this.state.id} />
+          );
+        }
+      }
+      const wrapper = mount(<Foo />);
+      expect(wrapper.state()).to.eql({ id: 'foo' });
+      expect(() => wrapper.setState({ id: 'bar' }, 1)).to.throw(
+        Error,
+        'setState(...): Expected the last optional `callback` argument ' +
+        'to be a function. Instead received: number.'
+      );
+    });
   });
 
   describe('.is(selector)', () => {

test/ShallowWrapper-spec.jsx

@@ -1078,7 +1078,7 @@ describe('shallow', () => {
     });
   });
 
-  describe('.setState(newState)', () => {
+  describe('.setState(newState[, callback])', () => {
     it('should set the state of the root node', () => {
       class Foo extends React.Component {
         constructor(props) {
@@ -1097,6 +1097,25 @@ describe('shallow', () => {
       expect(wrapper.find('.bar').length).to.equal(1);
     });
 
+    it('should call the callback when setState has completed', () => {
+      class Foo extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = { id: 'foo' };
+        }
+        render() {
+          return (
+            <div className={this.state.id} />
+          );
+        }
+      }
+      const wrapper = shallow(<Foo />);
+      expect(wrapper.state()).to.eql({ id: 'foo' });
+      wrapper.setState({ id: 'bar' }, () => {
+        expect(wrapper.state()).to.eql({ id: 'bar' });
+      });
+    });
+
     describeIf(!REACT013, 'stateless function components', () => {
       it('should throw when trying to access state', () => {
         const Foo = () => (