[Shallow]Call update() automatically through get querys

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

@@ -4650,13 +4650,16 @@ describe('shallow', () => {
       }
     }
 
-    it('should have updated output after an asynchronous setState', (done) => {
+    it('should have updated output after an asynchronous setState', () => {
       const wrapper = shallow(<Test />);
       wrapper.find('.async-btn').simulate('click');
-      setImmediate(() => {
-        wrapper.update();
+      return new Promise((resolve) => {
+        setImmediate(() => {
+          wrapper.update();
+          resolve();
+        });
+      }).then(() => {
         expect(wrapper.find('.show-me').length).to.equal(1);
-        done();
       });
     });
 
@@ -4692,4 +4695,86 @@ describe('shallow', () => {
       });
     });
   });
+
+  describe('setState through a props method', () => {
+    class Child extends React.Component {
+      render() {
+        return <button onClick={this.props.onClick}>click</button>;
+      }
+    }
+
+    it('should be able to get the latest state value', () => {
+      class App extends React.Component {
+        constructor(props) {
+          super(props);
+          this.state = {
+            count: 0,
+          };
+        }
+        onIncrement() {
+          this.setState({
+            count: this.state.count + 1,
+          });
+        }
+        render() {
+          return (
+            <div>
+              <Child onClick={() => this.onIncrement()} />
+              <p>{this.state.count}</p>
+            </div>
+          );
+        }
+      }
+      const wrapper = shallow(<App />);
+      const p = wrapper.find('p');
+      expect(wrapper.find('p').text()).to.equal('0');
+      wrapper.find(Child).prop('onClick')();
+      // this is still 0 because the wrapper won't be updated
+      expect(p.text()).to.equal('0');
+      expect(wrapper.find('p').text()).to.equal('1');
+    });
+  });
+
+  describe('setState through a props method in async', () => {
+    class Child extends React.Component {
+      render() {
+        return <button onClick={this.props.onClick}>click</button>;
+      }
+    }
+
+    it('should be able to get the latest state value', () => {
+      let App;
+      const promise = new Promise((resolve) => {
+        App = class extends React.Component {
+          constructor(props) {
+            super(props);
+            this.state = {
+              count: 0,
+            };
+          }
+          onIncrement() {
+            setTimeout(() => {
+              this.setState({
+                count: this.state.count + 1,
+              }, resolve);
+            });
+          }
+          render() {
+            return (
+              <div>
+                <Child onClick={() => this.onIncrement()} />
+                <p>{this.state.count}</p>
+              </div>
+            );
+          }
+        };
+      });
+      const wrapper = shallow(<App />);
+      expect(wrapper.find('p').text()).to.equal('0');
+      wrapper.find(Child).prop('onClick')();
+      return promise.then(() => {
+        expect(wrapper.find('p').text()).to.equal('1');
+      });
+    });
+  });
 });

packages/enzyme/src/ShallowWrapper.js

@@ -97,6 +97,10 @@ function getRootNode(node) {
   return node.rendered;
 }
 
+function getRootNodeInternal(wrapper) {
+  return wrapper[ROOT][NODE];
+}
+
 function privateSetNodes(wrapper, nodes) {
   if (!Array.isArray(nodes)) {
     privateSet(wrapper, NODE, nodes);
@@ -150,11 +154,13 @@ class ShallowWrapper {
     return this[ROOT];
   }
 
-
   getNodeInternal() {
     if (this.length !== 1) {
       throw new Error('ShallowWrapper::getNode() can only be called when wrapping one node');
     }
+    if (this[ROOT] === this) {
+      this.update();
+    }
     return this[NODE];
   }
 
@@ -167,7 +173,7 @@ class ShallowWrapper {
     if (this.length !== 1) {
       throw new Error('ShallowWrapper::getElement() can only be called when wrapping one node');
     }
-    return getAdapter(this[OPTIONS]).nodeToElement(this[NODE]);
+    return getAdapter(this[OPTIONS]).nodeToElement(this.getNodeInternal());
   }
 
   /**
@@ -176,7 +182,7 @@ class ShallowWrapper {
    * @return {Array<ReactElement>}
    */
   getElements() {
-    return this[NODES].map(getAdapter(this[OPTIONS]).nodeToElement);
+    return this.getNodesInternal().map(getAdapter(this[OPTIONS]).nodeToElement);
   }
 
   // eslint-disable-next-line class-methods-use-this
@@ -185,6 +191,9 @@ class ShallowWrapper {
   }
 
   getNodesInternal() {
+    if (this[ROOT] === this && this.length === 1) {
+      this.update();
+    }
     return this[NODES];
   }
 
@@ -225,9 +234,10 @@ class ShallowWrapper {
     if (this[ROOT] !== this) {
       throw new Error('ShallowWrapper::update() can only be called on the root');
     }
-    this.single('update', () => {
-      privateSetNodes(this, getRootNode(this[RENDERER].getNode()));
-    });
+    if (this.length !== 1) {
+      throw new Error('ShallowWrapper::update() can only be called when wrapping one node');
+    }
+    privateSetNodes(this, getRootNode(this[RENDERER].getNode()));
     return this;
   }
 
@@ -793,7 +803,7 @@ class ShallowWrapper {
    * @returns {ShallowWrapper}
    */
   parents(selector) {
-    const allParents = this.wrap(this.single('parents', n => parentsOfNode(n, this[ROOT][NODE])));
+    const allParents = this.wrap(this.single('parents', n => parentsOfNode(n, getRootNodeInternal(this))));
     return selector ? allParents.filter(selector) : allParents;
   }
 
@@ -1189,7 +1199,7 @@ if (ITERATOR_SYMBOL) {
   Object.defineProperty(ShallowWrapper.prototype, ITERATOR_SYMBOL, {
     configurable: true,
     value: function iterator() {
-      const iter = this[NODES][ITERATOR_SYMBOL]();
+      const iter = this.getNodesInternal()[ITERATOR_SYMBOL]();
       const adapter = getAdapter(this[OPTIONS]);
       return {
         next() {