From d7f63b9d5d1cc0e5d24bec06c6dffb2d0ac92e82 Mon Sep 17 00:00:00 2001
From: Jordan Harband <ljharb@gmail.com>
Date: Mon, 3 Sep 2018 00:43:21 -0700
Subject: [PATCH] [New] `mount`: `.state()`/`.setState()`: allow calling on
 children.

Fixes #635. Fixes #1289.
---
 .../test/ReactWrapper-spec.jsx                | 101 ++++++++++++++++++
 .../test/ShallowWrapper-spec.jsx              |  96 +++++++++++++++++
 packages/enzyme/src/ReactWrapper.js           |  10 +-
 3 files changed, 199 insertions(+), 8 deletions(-)

diff --git a/packages/enzyme-test-suite/test/ReactWrapper-spec.jsx b/packages/enzyme-test-suite/test/ReactWrapper-spec.jsx
index fd4f8c5e9..290a3e021 100644
--- a/packages/enzyme-test-suite/test/ReactWrapper-spec.jsx
+++ b/packages/enzyme-test-suite/test/ReactWrapper-spec.jsx
@@ -3045,6 +3045,63 @@ describeWithDOM('mount', () => {
 
       expect(mount(<Comp />).debug()).to.equal('<Comp />');
     });
+
+    describe('child components', () => {
+      class Child extends React.Component {
+        constructor(...args) {
+          super(...args);
+          this.state = { state: 'a' };
+        }
+
+        render() {
+          const { prop } = this.props;
+          const { state } = this.state;
+          return (
+            <div>
+              {prop} - {state}
+            </div>
+          );
+        }
+      }
+
+      class Parent extends React.Component {
+        constructor(...args) {
+          super(...args);
+          this.state = { childProp: 1 };
+        }
+
+        render() {
+          const { childProp } = this.state;
+          return <Child prop={childProp} />;
+        }
+      }
+
+      it('sets the state of the parent', () => {
+        const wrapper = mount(<Parent />);
+
+        expect(wrapper.text().trim()).to.eql('1 - a');
+
+        return new Promise((resolve) => {
+          wrapper.setState({ childProp: 2 }, () => {
+            expect(wrapper.text().trim()).to.eql('2 - a');
+            resolve();
+          });
+        });
+      });
+
+      it('sets the state of the child', () => {
+        const wrapper = mount(<Parent />);
+
+        expect(wrapper.text().trim()).to.eql('1 - a');
+
+        return new Promise((resolve) => {
+          wrapper.find(Child).setState({ state: 'b' }, () => {
+            expect(wrapper.text().trim()).to.eql('1 - b');
+            resolve();
+          });
+        });
+      });
+    });
   });
 
   describe('.is(selector)', () => {
@@ -3597,6 +3654,50 @@ describeWithDOM('mount', () => {
         expect(() => wrapper.state()).to.throw(Error, 'ReactWrapper::state() can only be called on class components');
       });
     });
+
+    describe('child components', () => {
+      class Child extends React.Component {
+        constructor(...args) {
+          super(...args);
+          this.state = { a: 'a' };
+        }
+
+        render() {
+          const { _ } = this.props;
+          const { a } = this.state;
+          return (
+            <div>
+              {_}
+              {a}
+            </div>
+          );
+        }
+      }
+
+      class Parent extends React.Component {
+        constructor(...args) {
+          super(...args);
+          this.state = { _: 1 };
+        }
+
+        render() {
+          const { _ } = this.state;
+          return <Child _={_} />;
+        }
+      }
+
+      it('gets the state of the parent', () => {
+        const wrapper = mount(<Parent />);
+
+        expect(wrapper.state()).to.eql({ _: 1 });
+      });
+
+      it('gets the state of the child', () => {
+        const wrapper = mount(<Parent />);
+
+        expect(wrapper.find(Child).state()).to.eql({ a: 'a' });
+      });
+    });
   });
 
   describe('.children([selector])', () => {
diff --git a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
index 13ce6c5b1..238ec3357 100644
--- a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
+++ b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
@@ -2738,6 +2738,59 @@ describe('shallow', () => {
 
       expect(shallow(<Comp />).debug()).to.equal('');
     });
+
+    describe('child components', () => {
+      class Child extends React.Component {
+        constructor(...args) {
+          super(...args);
+          this.state = { a: 'a' };
+        }
+
+        render() {
+          const { _ } = this.props;
+          const { a } = this.state;
+          return (
+            <div>
+              {_}
+              {a}
+            </div>
+          );
+        }
+      }
+
+      class Parent extends React.Component {
+        constructor(...args) {
+          super(...args);
+          this.state = { _: 1 };
+        }
+
+        render() {
+          const { _ } = this.state;
+          return <Child _={_} />;
+        }
+      }
+
+      it('sets the state of the parent', () => {
+        const wrapper = shallow(<Parent />);
+
+        expect(wrapper.debug()).to.eql('<Child _={1} />');
+
+        return new Promise((resolve) => {
+          wrapper.setState({ _: 2 }, () => {
+            expect(wrapper.debug()).to.eql('<Child _={2} />');
+            resolve();
+          });
+        });
+      });
+
+      it('can not set the state of the child', () => {
+        const wrapper = shallow(<Parent />);
+
+        expect(wrapper.debug()).to.eql('<Child _={1} />');
+
+        expect(() => wrapper.find(Child).setState({ a: 'b' })).to.throw(Error, 'ShallowWrapper::setState() can only be called on the root');
+      });
+    });
   });
 
   describe('.is(selector)', () => {
@@ -3291,6 +3344,49 @@ describe('shallow', () => {
         expect(() => wrapper.state()).to.throw(Error, 'ShallowWrapper::state() can only be called on class components');
       });
     });
+
+    describe('child components', () => {
+      class Child extends React.Component {
+        constructor(...args) {
+          super(...args);
+          this.state = { state: 'a' };
+        }
+
+        render() {
+          const { prop } = this.props;
+          const { state } = this.state;
+          return (
+            <div>
+              {prop} - {state}
+            </div>
+          );
+        }
+      }
+
+      class Parent extends React.Component {
+        constructor(...args) {
+          super(...args);
+          this.state = { childProp: 1 };
+        }
+
+        render() {
+          const { childProp } = this.state;
+          return <Child prop={childProp} />;
+        }
+      }
+
+      it('gets the state of the parent', () => {
+        const wrapper = shallow(<Parent />);
+
+        expect(wrapper.state()).to.eql({ childProp: 1 });
+      });
+
+      it('can not get the state of the child', () => {
+        const wrapper = shallow(<Parent />);
+
+        expect(() => wrapper.find(Child).state()).to.throw(Error, 'ShallowWrapper::state() can only be called on the root');
+      });
+    });
   });
 
   describe('.children([selector])', () => {
diff --git a/packages/enzyme/src/ReactWrapper.js b/packages/enzyme/src/ReactWrapper.js
index e3d577334..cc517fdd9 100644
--- a/packages/enzyme/src/ReactWrapper.js
+++ b/packages/enzyme/src/ReactWrapper.js
@@ -217,13 +217,13 @@ class ReactWrapper {
    * Forces a re-render. Useful to run before checking the render output if something external
    * may be updating the state of the component somewhere.
    *
-   * NOTE: can only be called on a wrapper instance that is also the root instance.
+   * NOTE: no matter what instance this is called on, it will always update the root.
    *
    * @returns {ReactWrapper}
    */
   update() {
     if (this[ROOT] !== this) {
-      throw new Error('ReactWrapper::update() can only be called on the root');
+      return this[ROOT].update();
     }
     privateSetNodes(this, this[RENDERER].getNode());
     return this;
@@ -307,9 +307,6 @@ class ReactWrapper {
    * @returns {ReactWrapper}
    */
   setState(state, callback = undefined) {
-    if (this[ROOT] !== this) {
-      throw new Error('ReactWrapper::setState() can only be called on the root');
-    }
     if (this.instance() === null || this[RENDERER].getNode().nodeType !== 'class') {
       throw new Error('ReactWrapper::setState() can only be called on class components');
     }
@@ -668,9 +665,6 @@ class ReactWrapper {
    * @returns {*}
    */
   state(name) {
-    if (this[ROOT] !== this) {
-      throw new Error('ReactWrapper::state() can only be called on the root');
-    }
     if (this.instance() === null || this[RENDERER].getNode().nodeType !== 'class') {
       throw new Error('ReactWrapper::state() can only be called on class components');
     }