From 813daabae048292aae7a0c415f6f7ed8a2401839 Mon Sep 17 00:00:00 2001
From: Toru Kobayashi <koba0004@gmail.com>
Date: Mon, 12 Feb 2018 18:02:48 +0900
Subject: [PATCH] ShallowWrapper calls update() automatically

---
 .../test/ShallowWrapper-spec.jsx              | 71 +++++++++++++++++++
 packages/enzyme/src/ShallowWrapper.js         | 26 +++++--
 2 files changed, 90 insertions(+), 7 deletions(-)

diff --git a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
index fc5da81d0..8d4d88b0e 100644
--- a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
+++ b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
@@ -4656,4 +4656,75 @@ describe('shallow', () => {
       });
     });
   });
+  describe('setState through a props method', () => {
+    it('should be able to get the latest state value', () => {
+      const Child = props => <button onClick={props.onClick}>click</button>;
+      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', () => {
+    it('should be able to get the latest state value', (done) => {
+      const Child = props => <button onClick={props.onClick}>click</button>;
+      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 />);
+      promise.then(() => {
+        expect(wrapper.find('p').text()).to.equal('1');
+        done();
+      });
+      expect(wrapper.find('p').text()).to.equal('0');
+      wrapper.find(Child).prop('onClick')();
+    });
+  });
 });
diff --git a/packages/enzyme/src/ShallowWrapper.js b/packages/enzyme/src/ShallowWrapper.js
index f28a8c30b..e9d3b648d 100644
--- a/packages/enzyme/src/ShallowWrapper.js
+++ b/packages/enzyme/src/ShallowWrapper.js
@@ -152,11 +152,17 @@ class ShallowWrapper {
     return this[ROOT];
   }
 
+  getRootNodeInternal() {
+    return this[ROOT][NODE];
+  }
 
   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];
   }
 
@@ -169,7 +175,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());
   }
 
   /**
@@ -178,7 +184,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
@@ -187,6 +193,9 @@ class ShallowWrapper {
   }
 
   getNodesInternal() {
+    if (this[ROOT] === this && this.length === 1) {
+      this.update();
+    }
     return this[NODES];
   }
 
@@ -227,9 +236,12 @@ 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');
+    }
+    // this.single('update', () => {
+    privateSetNodes(this, getRootNode(this[RENDERER].getNode()));
+    // });
     return this;
   }
 
@@ -795,7 +807,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, this.getRootNodeInternal())));
     return selector ? allParents.filter(selector) : allParents;
   }
 
@@ -1190,7 +1202,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() {