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

…paths as well. Fixes #1756
enzymejs · Aug 16, 2018 · e590784 · e590784
1 parent be087d4
commit e590784
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 10 deletions.
diff --git a/packages/enzyme-test-suite/test/ReactWrapper-spec.jsx b/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)', () => {

diff --git a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx b/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)', () => {

diff --git a/packages/enzyme/src/ShallowWrapper.js b/packages/enzyme/src/ShallowWrapper.js
@@ -163,21 +163,19 @@ 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 instance;
     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);
-      // Ensure to call componentDidUpdate when instance.setState is called
-      if (instance && lifecycles.componentDidUpdate.onSetState && !instance[SET_STATE]) {
-        privateSet(instance, SET_STATE, instance.setState);
-        instance.setState = (...args) => this.setState(...args);
-      }
+      ({ instance } = this[RENDERER].getNode());
       if (
         !options.disableLifecycleMethods
         && instance
@@ -195,6 +193,15 @@ class ShallowWrapper {
       privateSetNodes(this, nodes);
     }
     privateSet(this, OPTIONS, root ? root[OPTIONS] : options);
+
+    if (!instance) {
+      ({ instance } = this[RENDERER].getNode());
+    }
+    // Ensure to call componentDidUpdate when instance.setState is called
+    if (instance && lifecycles.componentDidUpdate.onSetState && !instance[SET_STATE]) {
+      privateSet(instance, SET_STATE, instance.setState);
+      instance.setState = (...args) => this.setState(...args);
+    }
   }
 
   /**