Fix to skip updates when nextState is null or undefined

Author: koba04

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

@@ -2460,6 +2460,64 @@ describeWithDOM('mount', () => {
       });
     });
 
+    it('should prevent the update if the nextState is null or undefined', () => {
+      class Foo extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = { id: 'foo' };
+        }
+
+        componentDidUpdate() {}
+
+        render() {
+          return (
+            <div className={this.state.id} />
+          );
+        }
+      }
+
+      const wrapper = mount(<Foo />);
+      const spy = sinon.spy(wrapper.instance(), 'componentDidUpdate');
+      const callback = sinon.spy();
+      wrapper.setState(() => ({ id: 'bar' }), callback);
+      expect(spy).to.have.property('callCount', 1);
+      expect(callback).to.have.property('callCount', 1);
+
+      wrapper.setState(() => null, callback);
+      expect(spy).to.have.property('callCount', 1);
+      // the callback should always be called
+      expect(callback).to.have.property('callCount', 2);
+
+      wrapper.setState(() => undefined, callback);
+      expect(spy).to.have.property('callCount', 1);
+      expect(callback).to.have.property('callCount', 3);
+    });
+
+    it('should prevent an infinite loop if the nextState is null or undefined from setState in CDU', () => {
+      class Foo extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = { id: 'foo' };
+        }
+
+        componentDidUpdate() {
+          // eslint-disable-next-line react/no-did-update-set-state
+          this.setState(() => null);
+        }
+
+        render() {
+          return (
+            <div className={this.state.id} />
+          );
+        }
+      }
+
+      const wrapper = mount(<Foo />);
+      const spy = sinon.spy(wrapper.instance(), 'componentDidUpdate');
+      wrapper.setState(() => ({ id: 'bar' }));
+      expect(spy).to.have.property('callCount', 1);
+    });
+
     describe('should not call componentWillReceiveProps after setState is called', () => {
       it('should not call componentWillReceiveProps upon rerender', () => {
         class A extends React.Component {

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

@@ -2406,6 +2406,64 @@ describe('shallow', () => {
       });
     });
 
+    it('should prevent the update if the nextState is null or undefined', () => {
+      class Foo extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = { id: 'foo' };
+        }
+
+        componentDidUpdate() {}
+
+        render() {
+          return (
+            <div className={this.state.id} />
+          );
+        }
+      }
+
+      const wrapper = shallow(<Foo />);
+      const spy = sinon.spy(wrapper.instance(), 'componentDidUpdate');
+      const callback = sinon.spy();
+      wrapper.setState(() => ({ id: 'bar' }), callback);
+      expect(spy).to.have.property('callCount', 1);
+      expect(callback).to.have.property('callCount', 1);
+
+      wrapper.setState(() => null, callback);
+      expect(spy).to.have.property('callCount', 1);
+      // the callback should always be called
+      expect(callback).to.have.property('callCount', 2);
+
+      wrapper.setState(() => undefined, callback);
+      expect(spy).to.have.property('callCount', 1);
+      expect(callback).to.have.property('callCount', 3);
+    });
+
+    it('should prevent an infinite loop if the nextState is null or undefined from setState in CDU', () => {
+      class Foo extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = { id: 'foo' };
+        }
+
+        componentDidUpdate() {
+          // eslint-disable-next-line react/no-did-update-set-state
+          this.setState(() => null);
+        }
+
+        render() {
+          return (
+            <div className={this.state.id} />
+          );
+        }
+      }
+
+      const wrapper = shallow(<Foo />);
+      const spy = sinon.spy(wrapper.instance(), 'componentDidUpdate');
+      wrapper.setState(() => ({ id: 'bar' }));
+      expect(spy).to.have.property('callCount', 1);
+    });
+
     describe('should not call componentWillReceiveProps after setState is called', () => {
       it('should not call componentWillReceiveProps upon rerender', () => {
         class A extends React.Component {

packages/enzyme/src/ShallowWrapper.js

@@ -436,6 +436,7 @@ class ShallowWrapper {
     if (arguments.length > 1 && typeof callback !== 'function') {
       throw new TypeError('ReactWrapper::setState() expects a function as its second argument');
     }
+
     this.single('setState', () => {
       withSetStateAllowed(() => {
         const adapter = getAdapter(this[OPTIONS]);
@@ -446,6 +447,14 @@ class ShallowWrapper {
         const prevProps = instance.props;
         const prevState = instance.state;
         const prevContext = instance.context;
+
+        if (typeof state === 'function') {
+          state = state.call(instance, prevState, prevProps);
+        }
+        // returning null or undefined prevents the update
+        // https://github.com/facebook/react/pull/12756
+        const hasUpdate = state !== null && state !== undefined;
+
         // When shouldComponentUpdate returns false we shouldn't call componentDidUpdate.
         // so we spy shouldComponentUpdate to get the result.
         let spy;
@@ -471,7 +480,8 @@ class ShallowWrapper {
           spy.restore();
         }
         if (
-          shouldRender
+          hasUpdate
+          && shouldRender
           && !this[OPTIONS].disableLifecycleMethods
           && lifecycles.componentDidUpdate
           && lifecycles.componentDidUpdate.onSetState