[Fix] shallow: skip updates when nextState is null or undefined

 - related: facebook/react#12756

Fixes #1783.

Author: koba04

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

@@ -2460,6 +2460,72 @@ describeWithDOM('mount', () => {
       });
     });
 
+    it('prevents the update if 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('prevents an infinite loop if 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} />
+          );
+        }
+      }
+
+      let payload;
+      const stub = sinon.stub(Foo.prototype, 'componentDidUpdate')
+        .callsFake(function componentDidUpdate() { this.setState(() => payload); });
+
+      const wrapper = mount(<Foo />);
+
+      wrapper.setState(() => ({ id: 'bar' }));
+      expect(stub).to.have.property('callCount', 1);
+
+      payload = null;
+      wrapper.setState(() => ({ id: 'bar' }));
+      expect(stub).to.have.property('callCount', 2);
+    });
+
     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,69 @@ describe('shallow', () => {
       });
     });
 
+    it('prevents the update if 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('prevents an infinite loop if nextState is null or undefined from setState in CDU', () => {
+      class Foo extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = { id: 'foo' };
+        }
+
+        componentDidUpdate() {}
+
+        render() {
+          return (
+            <div className={this.state.id} />
+          );
+        }
+      }
+
+      let payload;
+      const stub = sinon.stub(Foo.prototype, 'componentDidUpdate')
+        .callsFake(function componentDidUpdate() { this.setState(() => payload); });
+
+      const wrapper = shallow(<Foo />);
+
+      wrapper.setState(() => ({ id: 'bar' }));
+      expect(stub).to.have.property('callCount', 1);
+
+      payload = null;
+      wrapper.setState(() => ({ id: 'bar' }));
+      expect(stub).to.have.property('callCount', 2);
+    });
+
     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,15 @@ class ShallowWrapper {
         const prevProps = instance.props;
         const prevState = instance.state;
         const prevContext = instance.context;
+
+        const statePayload = typeof state === 'function'
+          ? state.call(instance, prevState, prevProps)
+          : state;
+
+        // returning null or undefined prevents the update
+        // https://github.com/facebook/react/pull/12756
+        const maybeHasUpdate = statePayload != null;
+
         // When shouldComponentUpdate returns false we shouldn't call componentDidUpdate.
         // so we spy shouldComponentUpdate to get the result.
         let spy;
@@ -462,16 +472,17 @@ class ShallowWrapper {
         // We don't pass the setState callback here
         // to guarantee to call the callback after finishing the render
         if (instance[SET_STATE]) {
-          instance[SET_STATE](state);
+          instance[SET_STATE](statePayload);
         } else {
-          instance.setState(state);
+          instance.setState(statePayload);
         }
         if (spy) {
           shouldRender = spy.getLastReturnValue();
           spy.restore();
         }
         if (
-          shouldRender
+          maybeHasUpdate
+          && shouldRender
           && !this[OPTIONS].disableLifecycleMethods
           && lifecycles.componentDidUpdate
           && lifecycles.componentDidUpdate.onSetState