Make setNodesInternal module-scoped; add basic ReactWrapper test

enzymejs · Jan 27, 2018 · 608a877 · 608a877
1 parent a5342e9
commit 608a877
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 28 deletions.
diff --git a/packages/enzyme-test-suite/test/ReactWrapper-spec.jsx b/packages/enzyme-test-suite/test/ReactWrapper-spec.jsx
@@ -3117,6 +3117,18 @@ describeWithDOM('mount', () => {
     expect(rendered.html()).to.equal(null);
   });
 
+  itIf(REACT16, 'works with class components that return arrays', () => {
+    class Foo extends React.Component {
+      render() {
+        return [<div />, <div />];
+      }
+    }
+    const wrapper = mount(<Foo />);
+    expect(wrapper).to.have.length(1);
+    expect(wrapper.type()).to.equal(Foo);
+    expect(wrapper.find('div')).to.have.length(2);
+  });
+
   itIf(is('>=15 || ^16.0.0-alpha'), 'works with SFCs that return null', () => {
     const Foo = () => null;
 

diff --git a/packages/enzyme/src/ReactWrapper.js b/packages/enzyme/src/ReactWrapper.js
@@ -63,6 +63,20 @@ function filterWhereUnwrapped(wrapper, predicate) {
   return wrapper.wrap(compact(wrapper.getNodesInternal().filter(predicate)));
 }
 
+function privateSetNodes(wrapper, nodes) {
+  if (!nodes) {
+    privateSet(wrapper, NODE, null);
+    privateSet(wrapper, NODES, []);
+  } else if (!Array.isArray(nodes)) {
+    privateSet(wrapper, NODE, nodes);
+    privateSet(wrapper, NODES, [nodes]);
+  } else {
+    privateSet(wrapper, NODE, nodes[0]);
+    privateSet(wrapper, NODES, nodes);
+  }
+  privateSet(wrapper, 'length', wrapper[NODES].length);
+}
+
 /**
  * @class ReactWrapper
  */
@@ -79,25 +93,12 @@ class ReactWrapper {
       privateSet(this, RENDERER, renderer);
       renderer.render(nodes, options.context);
       privateSet(this, ROOT, this);
-      const node = this[RENDERER].getNode();
-      privateSet(this, NODE, node);
-      privateSet(this, NODES, [node]);
-      this.length = 1;
+      privateSetNodes(this, this[RENDERER].getNode());
     } else {
       privateSet(this, UNRENDERED, null);
       privateSet(this, RENDERER, root[RENDERER]);
       privateSet(this, ROOT, root);
-      if (!nodes) {
-        privateSet(this, NODE, null);
-        privateSet(this, NODES, []);
-      } else if (!Array.isArray(nodes)) {
-        privateSet(this, NODE, nodes);
-        privateSet(this, NODES, [nodes]);
-      } else {
-        privateSet(this, NODE, nodes[0]);
-        privateSet(this, NODES, nodes);
-      }
-      this.length = this[NODES].length;
+      privateSetNodes(this, nodes);
     }
     privateSet(this, OPTIONS, root ? root[OPTIONS] : options);
   }

diff --git a/packages/enzyme/src/ShallowWrapper.js b/packages/enzyme/src/ShallowWrapper.js
@@ -99,6 +99,17 @@ function getRootNode(node) {
   return node.rendered;
 }
 
+function privateSetNodes(wrapper, nodes) {
+  if (!Array.isArray(nodes)) {
+    privateSet(wrapper, NODE, nodes);
+    privateSet(wrapper, NODES, [nodes]);
+  } else {
+    privateSet(wrapper, NODE, nodes[0]);
+    privateSet(wrapper, NODES, nodes);
+  }
+  privateSet(wrapper, 'length', wrapper[NODES].length);
+}
+
 /**
  * @class ShallowWrapper
  */
@@ -122,27 +133,16 @@ class ShallowWrapper {
           instance.componentDidMount();
         });
       }
-      this.setNodesInternal(getRootNode(this[RENDERER].getNode()));
+      privateSetNodes(this, getRootNode(this[RENDERER].getNode()));
     } else {
       privateSet(this, ROOT, root);
       privateSet(this, UNRENDERED, null);
       privateSet(this, RENDERER, root[RENDERER]);
-      this.setNodesInternal(nodes);
+      privateSetNodes(this, nodes);
     }
     privateSet(this, OPTIONS, root ? root[OPTIONS] : options);
   }
 
-  setNodesInternal(nodes) {
-    if (!Array.isArray(nodes)) {
-      privateSet(this, NODE, nodes);
-      privateSet(this, NODES, [nodes]);
-    } else {
-      privateSet(this, NODE, nodes[0]);
-      privateSet(this, NODES, nodes);
-    }
-    this.length = this[NODES].length;
-  }
-
   /**
    * Returns the root wrapper
    *