don't need to track initialNode separately

enzymejs · Oct 24, 2016 · bcf175c · bcf175c
1 parent 9172828
commit bcf175c
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git 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
@@ -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();
   }
 
   /**