don't need to track initialNode separately

enzymejs · Sep 17, 2016 · 663a552 · 663a552
1 parent 79ac4a7
commit 663a552
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
@@ -89,10 +89,8 @@ export default 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;
@@ -111,7 +109,12 @@ export default 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
@@ -87,10 +87,8 @@ export default 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;
@@ -105,7 +103,12 @@ export default 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();
   }
 
   /**