From 34d69ce74b79acfa7edcc2f7f7f379d3cce6a6bc Mon Sep 17 00:00:00 2001
From: Emilien Muraton <emilien.muraton@gmail.com>
Date: Sun, 22 Jul 2018 01:14:58 +0100
Subject: [PATCH] [New] `shallow`: `setProps()`: Add callback argument

---
 .../enzyme-test-suite/test/ShallowWrapper-spec.jsx   |  2 +-
 packages/enzyme/src/ShallowWrapper.js                | 12 ++++++++++--
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
index 1e4885bf3..ca1c3647a 100644
--- a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
+++ b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
@@ -1296,8 +1296,8 @@ describe('shallow', () => {
         wrapper.setProps({ id: 'bar', foo: 'bla' }, () => {
           expect(wrapper.find('.bar')).to.have.lengthOf(1);
         });
-        expect(wrapper.find('.bar')).to.have.lengthOf(0);
       });
+      expect(wrapper.find('.foo')).to.have.lengthOf(0);
     });
 
     it('should call componentWillReceiveProps, shouldComponentUpdate, componentWillUpdate, and componentDidUpdate with merged newProps', () => {
diff --git a/packages/enzyme/src/ShallowWrapper.js b/packages/enzyme/src/ShallowWrapper.js
index c6602a514..edba7dcd1 100644
--- a/packages/enzyme/src/ShallowWrapper.js
+++ b/packages/enzyme/src/ShallowWrapper.js
@@ -30,6 +30,8 @@ import {
 } from './RSTTraversal';
 import { buildPredicate, reduceTreesBySelector } from './selectors';
 
+const noop = () => {};
+
 const NODE = sym('__node__');
 const NODES = sym('__nodes__');
 const RENDERER = sym('__renderer__');
@@ -376,13 +378,19 @@ class ShallowWrapper {
    * 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 {ShallowWrapper}
    */
-  setProps(props) {
+  setProps(props, callback = noop) {
     if (this[ROOT] !== this) {
       throw new Error('ShallowWrapper::setProps() can only be called on the root');
     }
-    return this.rerender(props);
+    if (typeof callback !== 'function') {
+      throw new TypeError('ShallowWrapper::setProps() expects a function as its second argument');
+    }
+    this.rerender(props);
+    callback();
+    return this;
   }
 
   /**