Fix to call a setState callback after finishing the the render

enzymejs · Dec 26, 2017 · 04076f6 · 04076f6
1 parent c8f828a
commit 04076f6
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
@@ -1378,8 +1378,10 @@ describe('shallow', () => {
       }
       const wrapper = shallow(<Foo />);
       expect(wrapper.state()).to.eql({ id: 'foo' });
-      wrapper.setState({ id: 'bar' }, () => {
+      wrapper.setState({ id: 'bar' }, function callback() {
         expect(wrapper.state()).to.eql({ id: 'bar' });
+        expect(this.state).to.eql({ id: 'bar' });
+        expect(wrapper.find('div').prop('className')).to.eql('bar');
       });
     });
 

diff --git a/packages/enzyme/src/ShallowWrapper.js b/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) {
+          callback.call(instance);
+        }
       });
     });
     return this;