[New] mount: .state()/.setState(): allow calling on children

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,49 @@ 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 = { 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 = mount(<Parent />);
+
+        expect(wrapper.state()).to.eql({ childProp: 1 });
+      });
+
+      it('gets the state of the child', () => {
+        const wrapper = mount(<Parent />);
+
+        expect(wrapper.find(Child).state()).to.eql({ state: 'a' });
+      });
+    });
   });
 
   describe('.children([selector])', () => {

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

@@ -2738,6 +2738,61 @@ describe('shallow', () => {
 
       expect(shallow(<Comp />).debug()).to.equal('');
     });
+
+    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 = shallow(<Parent />);
+
+        expect(wrapper.debug()).to.eql('<Child prop={1} />');
+
+        return new Promise((resolve) => {
+          wrapper.setState({ childProp: 2 }, () => {
+            expect(wrapper.debug()).to.eql('<Child prop={2} />');
+            resolve();
+          });
+        });
+      });
+
+      it('can not set the state of the child', () => {
+        const wrapper = shallow(<Parent />);
+
+        expect(wrapper.debug()).to.eql('<Child prop={1} />');
+
+        expect(() => wrapper.find(Child).setState({ state: 'b' })).to.throw(
+          Error,
+          'ShallowWrapper::setState() can only be called on the root',
+        );
+      });
+    });
   });
 
   describe('.is(selector)', () => {
@@ -3291,6 +3346,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])', () => {

packages/enzyme/src/ReactWrapper.js

@@ -223,13 +223,14 @@ 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');
+    const root = this[ROOT];
+    if (this !== root) {
+      return root.update();
     }
     privateSetNodes(this, this[RENDERER].getNode());
     return this;
@@ -313,9 +314,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');
     }
@@ -675,9 +673,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');
     }