Fix to call a setState callback after finishing the the render

packages/enzyme-adapter-react-15.4/package.json

@@ -34,6 +34,7 @@
   "author": "Leland Richardson <leland.richardson@airbnb.com>",
   "license": "MIT",
   "dependencies": {
+    "enzyme-adapter-react-helper": "^1.2.2",
     "enzyme-adapter-utils": "^1.3.0",
     "lodash": "^4.17.4",
     "object.assign": "^4.1.0",

packages/enzyme-adapter-react-15.4/src/ReactFifteenFourAdapter.js

@@ -16,6 +16,7 @@ import {
   createMountWrapper,
   propsWithKeysAndRef,
 } from 'enzyme-adapter-utils';
+import ifReact from 'enzyme-adapter-react-helper/build/ifReact';
 
 function compositeTypeToNodeType(type) {
   switch (type) {
@@ -252,6 +253,15 @@ class ReactFifteenFourAdapter extends EnzymeAdapter {
   createElement(...args) {
     return React.createElement(...args);
   }
+
+  invokeSetStateCallback(instance, callback) {
+    // React in >= 15.4, and < 16 pass undefined to a setState callback
+    ifReact(
+      '^15.4',
+      () => { callback.call(instance, undefined); },
+      () => { super.invokeSetStateCallback(instance, callback); },
+    );
+  }
 }
 
 module.exports = ReactFifteenFourAdapter;

packages/enzyme-adapter-react-15/src/ReactFifteenAdapter.js

@@ -253,6 +253,11 @@ class ReactFifteenAdapter extends EnzymeAdapter {
   createElement(...args) {
     return React.createElement(...args);
   }
+
+  invokeSetStateCallback(instance, callback) {
+    // React in >= 15.4, and < 16 pass undefined to a setState callback
+    callback.call(instance, undefined);
+  }
 }
 
 module.exports = ReactFifteenAdapter;

packages/enzyme-test-suite/package.json

@@ -49,4 +49,4 @@
     "eslint-plugin-jsx-a11y": "^6.0.3",
     "eslint-plugin-react": "^7.5.1"
   }
-}
+}

packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx

@@ -8,12 +8,15 @@ import { ITERATOR_SYMBOL, withSetStateAllowed, sym } from 'enzyme/build/Utils';
 import './_helpers/setupAdapters';
 import { createClass } from './_helpers/react-compat';
 import { describeIf, itIf, itWithData, generateEmptyRenderData } from './_helpers';
-import { REACT013, REACT014, REACT16, is } from './_helpers/version';
+import { REACT013, REACT014, REACT15, REACT150_4, REACT16, is } from './_helpers/version';
 
 // The shallow renderer in react 16 does not yet support batched updates. When it does,
 // we should be able to go un-skip all of the tests that are skipped with this flag.
 const BATCHING = !REACT16;
 
+// some React versions pass undefined as an argument of setState callback.
+const CALLING_SETSTATE_CALLBACK_WITH_UNDEFINED = REACT15 && !REACT150_4;
+
 const getElementPropSelector = prop => x => x.props[prop];
 const getWrapperPropSelector = prop => x => x.prop(prop);
 
@@ -1378,8 +1381,11 @@ describe('shallow', () => {
       }
       const wrapper = shallow(<Foo />);
       expect(wrapper.state()).to.eql({ id: 'foo' });
-      wrapper.setState({ id: 'bar' }, () => {
+      wrapper.setState({ id: 'bar' }, function callback(...args) {
         expect(wrapper.state()).to.eql({ id: 'bar' });
+        expect(this.state).to.eql({ id: 'bar' });
+        expect(wrapper.find('div').prop('className')).to.eql('bar');
+        expect(args).to.eql(CALLING_SETSTATE_CALLBACK_WITH_UNDEFINED ? [undefined] : []);
       });
     });
 

packages/enzyme-test-suite/test/_helpers/setupAdapters.js

@@ -5,20 +5,26 @@
  * version of React is loaded, and configures enzyme to use the right
  * corresponding adapter.
  */
-const Version = require('./version');
+const {
+  REACT013,
+  REACT014,
+  REACT15,
+  REACT155,
+  REACT16,
+} = require('./version');
 const Enzyme = require('enzyme');
 
 let Adapter = null;
 
-if (Version.REACT013) {
+if (REACT013) {
   Adapter = require('enzyme-adapter-react-13');
-} else if (Version.REACT014) {
+} else if (REACT014) {
   Adapter = require('enzyme-adapter-react-14');
-} else if (Version.REACT155) {
+} else if (REACT155) {
   Adapter = require('enzyme-adapter-react-15');
-} else if (Version.REACT15) {
+} else if (REACT15) {
   Adapter = require('enzyme-adapter-react-15.4');
-} else if (Version.REACT16) {
+} else if (REACT16) {
   Adapter = require('enzyme-adapter-react-16');
 }
 

packages/enzyme-test-suite/test/_helpers/version.js

@@ -8,6 +8,7 @@ const [major, minor] = VERSION.split('.');
 export const REACT013 = VERSION.slice(0, 4) === '0.13';
 export const REACT014 = VERSION.slice(0, 4) === '0.14';
 export const REACT15 = major === '15';
+export const REACT150_4 = REACT15 && minor < 5;
 export const REACT155 = REACT15 && minor >= 5;
 export const REACT16 = major === '16';
 

packages/enzyme/src/EnzymeAdapter.js

@@ -30,6 +30,11 @@ class EnzymeAdapter {
   createElement(type, props, ...children) {
     throw unimplementedError('createElement', 'EnzymeAdapter');
   }
+
+  // eslint-disable-next-line class-methods-use-this
+  invokeSetStateCallback(instance, callback) {
+    callback.call(instance);
+  }
 }
 
 EnzymeAdapter.MODES = {

packages/enzyme/src/ShallowWrapper.js

@@ -385,7 +385,9 @@ class ShallowWrapper {
             return shouldRender;
           };
         }
-        instance.setState(state, callback);
+        // We don't pass the setState callback here
+        // to guarantee to call the callback after finishing the render
+        instance.setState(state);
         if (
           shouldRender &&
           !this[OPTIONS].disableLifecycleMethods &&
@@ -400,6 +402,10 @@ class ShallowWrapper {
           }
         }
         this.update();
+        // call the setState callback
+        if (callback) {
+          adapter.invokeSetStateCallback(instance, callback);
+        }
       });
     });
     return this;