[Fix] shallow: .setState(): stub out setState on non-root code paths as well.

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

@@ -2317,6 +2317,31 @@ describeWithDOM('mount', () => {
       expect(wrapper.state()).to.eql({ id: 'foo' });
       expect(() => wrapper.setState({ id: 'bar' }, 1)).to.throw(Error);
     });
+
+    it('should preserve the receiver', () => {
+      class Comp extends React.Component {
+        constructor(...args) {
+          super(...args);
+
+          this.state = {
+            key: '',
+          };
+
+          this.instanceFunction = () => this.setState(() => ({ key: 'value' }));
+        }
+
+        componentDidMount() {
+          this.instanceFunction();
+        }
+
+        render() {
+          const { key } = this.state;
+          return key ? null : null;
+        }
+      }
+
+      expect(mount(<Comp />).debug()).to.equal('<Comp />');
+    });
   });
 
   describe('.is(selector)', () => {
@@ -5144,5 +5169,32 @@ describeWithDOM('mount', () => {
       expect(wrapper.state('foo')).to.equal('onChange update');
       expect(spy).to.have.property('callCount', 1);
     });
+
+    it('should call `componentDidUpdate` when component’s `setState` is called', () => {
+      class Foo extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = {
+            foo: 'init',
+          };
+          this.update = () => this.setState({ foo: 'update' });
+        }
+
+        componentDidMount() {
+          this.update();
+        }
+
+        componentDidUpdate() {}
+
+        render() {
+          return <div>{this.state.foo}</div>;
+        }
+      }
+      const spy = sinon.spy(Foo.prototype, 'componentDidUpdate');
+
+      const wrapper = mount(<Foo />);
+      expect(spy).to.have.property('callCount', 1);
+      expect(wrapper.state('foo')).to.equal('update');
+    });
   });
 });

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

@@ -2142,7 +2142,7 @@ describe('shallow', () => {
       expect(wrapper.find('.bar')).to.have.lengthOf(1);
     });
 
-    it.skip('allows setState inside of componentDidMount', () => {
+    it('allows setState inside of componentDidMount', () => {
       class MySharona extends React.Component {
         constructor(props) {
           super(props);
@@ -2233,6 +2233,31 @@ describe('shallow', () => {
       expect(wrapper.state()).to.eql({ id: 'foo' });
       expect(() => wrapper.setState({ id: 'bar' }, 1)).to.throw(Error);
     });
+
+    it('should preserve the receiver', () => {
+      class Comp extends React.Component {
+        constructor(...args) {
+          super(...args);
+
+          this.state = {
+            key: '',
+          };
+
+          this.instanceFunction = () => this.setState(() => ({ key: 'value' }));
+        }
+
+        componentDidMount() {
+          this.instanceFunction();
+        }
+
+        render() {
+          const { key } = this.state;
+          return key ? null : null;
+        }
+      }
+
+      expect(shallow(<Comp />).debug()).to.equal('');
+    });
   });
 
   describe('.is(selector)', () => {
@@ -4942,6 +4967,33 @@ describe('shallow', () => {
         expect(wrapper.state('foo')).to.equal('onChange update');
         expect(spy).to.have.property('callCount', 1);
       });
+
+      it('should call `componentDidUpdate` when component’s `setState` is called', () => {
+        class Foo extends React.Component {
+          constructor(props) {
+            super(props);
+            this.state = {
+              foo: 'init',
+            };
+            this.update = () => this.setState({ foo: 'update' });
+          }
+
+          componentDidMount() {
+            this.update();
+          }
+
+          componentDidUpdate() {}
+
+          render() {
+            return <div>{this.state.foo}</div>;
+          }
+        }
+        const spy = sinon.spy(Foo.prototype, 'componentDidUpdate');
+
+        const wrapper = shallow(<Foo />);
+        expect(spy).to.have.property('callCount', 1);
+        expect(wrapper.state('foo')).to.equal('update');
+      });
     });
 
     describeIf(is('>= 16'), 'support getSnapshotBeforeUpdate', () => {

packages/enzyme/src/ShallowWrapper.js

@@ -163,38 +163,43 @@ function privateSetNodes(wrapper, nodes) {
 class ShallowWrapper {
   constructor(nodes, root, passedOptions = {}) {
     validateOptions(passedOptions);
+
     const options = makeOptions(passedOptions);
+    const adapter = getAdapter(options);
+    const lifecycles = getAdapterLifecycles(adapter);
+
+    let renderedNode;
     if (!root) {
       privateSet(this, ROOT, this);
       privateSet(this, UNRENDERED, nodes);
-      const renderer = getAdapter(options).createRenderer({ mode: 'shallow', ...options });
+      const renderer = adapter.createRenderer({ mode: 'shallow', ...options });
       privateSet(this, RENDERER, renderer);
       this[RENDERER].render(nodes, options.context);
-      const { instance } = this[RENDERER].getNode();
-      const adapter = getAdapter(this[OPTIONS]);
-      const lifecycles = getAdapterLifecycles(adapter);
+      renderedNode = this[RENDERER].getNode();
+      privateSetNodes(this, getRootNode(renderedNode));
+    } else {
+      privateSet(this, ROOT, root);
+      privateSet(this, UNRENDERED, null);
+      privateSet(this, RENDERER, root[RENDERER]);
+      privateSetNodes(this, nodes);
+      renderedNode = this[RENDERER].getNode();
+    }
+    privateSet(this, OPTIONS, root ? root[OPTIONS] : options);
+
+    const { instance } = renderedNode;
+    if (instance && !options.disableLifecycleMethods) {
       // Ensure to call componentDidUpdate when instance.setState is called
-      if (instance && lifecycles.componentDidUpdate.onSetState && !instance[SET_STATE]) {
+      if (lifecycles.componentDidUpdate.onSetState && !instance[SET_STATE]) {
         privateSet(instance, SET_STATE, instance.setState);
         instance.setState = (...args) => this.setState(...args);
       }
-      if (
-        !options.disableLifecycleMethods
-        && instance
-        && typeof instance.componentDidMount === 'function'
-      ) {
+
+      if (typeof instance.componentDidMount === 'function') {
         this[RENDERER].batchedUpdates(() => {
           instance.componentDidMount();
         });
       }
-      privateSetNodes(this, getRootNode(this[RENDERER].getNode()));
-    } else {
-      privateSet(this, ROOT, root);
-      privateSet(this, UNRENDERED, null);
-      privateSet(this, RENDERER, root[RENDERER]);
-      privateSetNodes(this, nodes);
     }
-    privateSet(this, OPTIONS, root ? root[OPTIONS] : options);
   }
 
   /**