From bcf175c907fcd279450709e1ad74592d1e541b2f Mon Sep 17 00:00:00 2001
From: Justin Bay <justin.bay@outlook.com>
Date: Sat, 17 Sep 2016 00:39:44 -0400
Subject: [PATCH] don't need to track initialNode separately

---
 src/ReactWrapper.jsx  | 9 ++++++---
 src/ShallowWrapper.js | 9 ++++++---
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/ReactWrapper.jsx b/src/ReactWrapper.jsx
index 11da3dda7..2f22a7b22 100644
--- a/src/ReactWrapper.jsx
+++ b/src/ReactWrapper.jsx
@@ -90,10 +90,8 @@ class ReactWrapper {
       if (!nodes) {
         this.initialNodes = [];
       } else if (!Array.isArray(nodes)) {
-        this.initialNode = nodes;
         this.initialNodes = [nodes];
       } else {
-        this.initialNode = nodes[0];
         this.initialNodes = nodes;
       }
       this.length = this.initialNodes.length;
@@ -112,7 +110,12 @@ class ReactWrapper {
    * @return {ReactComponent}
    */
   getNode() {
-    return this.initialNode || this.wrappedComponent;
+    if (this.length !== 1) {
+      throw new Error(
+        'ReactWrapper::getNode() can only be called when wrapping one node'
+      );
+    }
+    return this.initialNodes ? this.initialNodes[0] : this.wrappedComponent;
   }
 
   /**
diff --git a/src/ShallowWrapper.js b/src/ShallowWrapper.js
index 7e731c1fb..c07e8e5d5 100644
--- a/src/ShallowWrapper.js
+++ b/src/ShallowWrapper.js
@@ -90,10 +90,8 @@ class ShallowWrapper {
       this.unrendered = null;
       this.renderer = null;
       if (!Array.isArray(nodes)) {
-        this.initialNode = nodes;
         this.initialNodes = [nodes];
       } else {
-        this.initialNode = nodes[0];
         this.initialNodes = nodes;
       }
       this.length = this.initialNodes.length;
@@ -108,7 +106,12 @@ class ShallowWrapper {
    * @return {ReactElement}
    */
   getNode() {
-    return this.initialNode || this.renderer.getRenderOutput();
+    if (this.length !== 1) {
+      throw new Error(
+        'ShallowWrapper::getNode() can only be called when wrapping one node'
+      );
+    }
+    return this.initialNodes ? this.initialNodes[0] : this.renderer.getRenderOutput();
   }
 
   /**