From 09ba714ea47ead3415335ada6a12fd454ec2d39e Mon Sep 17 00:00:00 2001
From: Justin Bay <justin.bay@outlook.com>
Date: Wed, 27 Sep 2017 20:05:49 -0400
Subject: [PATCH] auto update wrappers behind a flag

---
 .../test/ReactWrapper-spec.jsx                | 49 +++++++++++++++++++
 .../test/ShallowWrapper-spec.jsx              |  8 ++-
 packages/enzyme/src/ReactWrapper.js           | 10 +++-
 packages/enzyme/src/ShallowWrapper.js         |  8 +++
 4 files changed, 69 insertions(+), 6 deletions(-)

diff --git a/packages/enzyme-test-suite/test/ReactWrapper-spec.jsx b/packages/enzyme-test-suite/test/ReactWrapper-spec.jsx
index 25f2a326c..40acc6daa 100644
--- a/packages/enzyme-test-suite/test/ReactWrapper-spec.jsx
+++ b/packages/enzyme-test-suite/test/ReactWrapper-spec.jsx
@@ -3683,6 +3683,55 @@ describeWithDOM('mount', () => {
     });
   });
 
+  describe('out-of-band state updates with autoUpdate', () => {
+    class Child extends React.Component {
+      render() {
+        return <span />;
+      }
+    }
+
+    class Test extends React.Component {
+      componentWillMount() {
+        this.state = {};
+      }
+
+      asyncSetState() {
+        setImmediate(() => {
+          this.setState({ showSpan: true });
+        });
+      }
+
+      callbackSetState() {
+        this.setState({ showSpan: true });
+      }
+
+      render() {
+        return (
+          <div>
+            {this.state && this.state.showSpan && <span className="show-me" />}
+            <button className="async-btn" onClick={() => this.asyncSetState()} />
+            <Child callback={() => this.callbackSetState()} />
+          </div>
+        );
+      }
+    }
+
+    it('should have updated output after an asynchronous setState', (done) => {
+      const wrapper = mount(<Test />, { autoUpdate: true });
+      wrapper.find('.async-btn').simulate('click');
+      setImmediate(() => {
+        expect(wrapper.find('.show-me').length).to.equal(1);
+        done();
+      });
+    });
+
+    it('should have updated output after child prop callback invokes setState', () => {
+      const wrapper = mount(<Test />, { autoUpdate: true });
+      wrapper.find(Child).props().callback();
+      expect(wrapper.find('.show-me').length).to.equal(1);
+    });
+  });
+
   describe('#single()', () => {
     it('throws if run on multiple nodes', () => {
       const wrapper = mount(<div><i /><i /></div>).children();
diff --git a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
index a724d8b6b..027475359 100644
--- a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
+++ b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
@@ -4446,7 +4446,7 @@ describe('shallow', () => {
     });
   });
 
-  describe('out-of-band state updates', () => {
+  describe('out-of-band state updates with autoUpdate', () => {
     class Child extends React.Component {
       render() {
         return <span />;
@@ -4486,19 +4486,17 @@ describe('shallow', () => {
     }
 
     it('should have updated output after an asynchronous setState', (done) => {
-      const wrapper = shallow(<Test />);
+      const wrapper = shallow(<Test />, { autoUpdate: true });
       wrapper.find('.async-btn').simulate('click');
       setImmediate(() => {
-        wrapper.update();
         expect(wrapper.find('.show-me').length).to.equal(1);
         done();
       });
     });
 
     it('should have updated output after child prop callback invokes setState', () => {
-      const wrapper = shallow(<Test />);
+      const wrapper = shallow(<Test />, { autoUpdate: true });
       wrapper.find(Child).props().callback();
-      wrapper.update();
       expect(wrapper.find('.show-me').length).to.equal(1);
     });
   });
diff --git a/packages/enzyme/src/ReactWrapper.js b/packages/enzyme/src/ReactWrapper.js
index b967b3cb8..f16132e36 100644
--- a/packages/enzyme/src/ReactWrapper.js
+++ b/packages/enzyme/src/ReactWrapper.js
@@ -124,6 +124,10 @@ class ReactWrapper {
         'ReactWrapper::getNode() can only be called when wrapping one node',
       );
     }
+    if (this[ROOT] === this && this[OPTIONS].autoUpdate) {
+      this[NODE] = this[RENDERER].getNode();
+      this[NODES] = [this[NODE]];
+    }
     return this[NODES][0];
   }
 
@@ -133,6 +137,10 @@ class ReactWrapper {
    * @return {Array<ReactComponent>}
    */
   getNodesInternal() {
+    if (this[ROOT] === this && this[OPTIONS].autoUpdate) {
+      this[NODE] = this[RENDERER].getNode();
+      this[NODES] = [this[NODE]];
+    }
     return this[NODES];
   }
 
@@ -675,7 +683,7 @@ class ReactWrapper {
    */
   parents(selector) {
     const allParents = this.wrap(
-      this.single('parents', n => parentsOfNode(n, this[ROOT].getNodeInternal())),
+      this.single('parents', n => parentsOfNode(n, this[ROOT][NODE])),
     );
     return selector ? allParents.filter(selector) : allParents;
   }
diff --git a/packages/enzyme/src/ShallowWrapper.js b/packages/enzyme/src/ShallowWrapper.js
index f1c207109..5bb8ed9f7 100644
--- a/packages/enzyme/src/ShallowWrapper.js
+++ b/packages/enzyme/src/ShallowWrapper.js
@@ -164,6 +164,10 @@ class ShallowWrapper {
         'ShallowWrapper::getNode() can only be called when wrapping one node',
       );
     }
+    if (this[ROOT] === this && this[OPTIONS].autoUpdate) {
+      this[NODE] = getRootNode(this[RENDERER].getNode());
+      this[NODES] = [this[NODE]];
+    }
     return this[NODE];
   }
 
@@ -198,6 +202,10 @@ class ShallowWrapper {
   }
 
   getNodesInternal() {
+    if (this[ROOT] === this && this[OPTIONS].autoUpdate) {
+      this[NODE] = getRootNode(this[RENDERER].getNode());
+      this[NODES] = [this[NODE]];
+    }
     return this[NODES];
   }