Fix ShallowWrapper for array-rendering components

packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js

@@ -231,7 +231,9 @@ class ReactSixteenAdapter extends EnzymeAdapter {
           key: cachedNode.key || undefined,
           ref: cachedNode.ref,
           instance: renderer._instance,
-          rendered: elementToTree(output),
+          rendered: Array.isArray(output)
+            ? flatten(output).map(elementToTree)
+            : elementToTree(output),
         };
       },
       simulateEvent(node, event, ...args) {

packages/enzyme-test-suite/test/ReactWrapper-spec.jsx

@@ -3117,6 +3117,19 @@ 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.lengthOf(1);
+    expect(wrapper.type()).to.equal(Foo);
+    expect(wrapper.children()).to.have.lengthOf(2);
+    expect(wrapper.find('div')).to.have.lengthOf(2);
+  });
+
   itIf(is('>=15 || ^16.0.0-alpha'), 'works with SFCs that return null', () => {
     const Foo = () => null;
 

packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx

@@ -3964,6 +3964,17 @@ describe('shallow', () => {
     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 = shallow(<Foo />);
+    expect(wrapper).to.have.lengthOf(2);
+    expect(wrapper.find('div')).to.have.lengthOf(2);
+  });
+
   itIf(is('>=15 || ^16.0.0-alpha'), 'works with SFCs that return null', () => {
     const Foo = () => null;
 

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);
   }
@@ -223,11 +224,7 @@ class ReactWrapper {
     if (this[ROOT] !== this) {
       throw new Error('ReactWrapper::update() can only be called on the root');
     }
-    this.single('update', () => {
-      const node = this[RENDERER].getNode();
-      this[NODE] = node;
-      this[NODES] = [node];
-    });
+    privateSetNodes(this, this[RENDERER].getNode());
     return this;
   }
 
@@ -258,9 +255,7 @@ class ReactWrapper {
     if (this[ROOT] !== this) {
       throw new Error('ReactWrapper::mount() can only be called on the root');
     }
-    this.single('mount', () => {
-      this[RENDERER].render(this[UNRENDERED], this[OPTIONS].context, () => this.update());
-    });
+    this[RENDERER].render(this[UNRENDERED], this[OPTIONS].context, () => this.update());
     return this;
   }
 

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,26 +133,16 @@ class ShallowWrapper {
           instance.componentDidMount();
         });
       }
-      privateSet(this, NODE, getRootNode(this[RENDERER].getNode()));
-      privateSet(this, NODES, [this[NODE]]);
-      this.length = 1;
+      privateSetNodes(this, getRootNode(this[RENDERER].getNode()));
     } else {
       privateSet(this, ROOT, root);
       privateSet(this, UNRENDERED, null);
       privateSet(this, RENDERER, root[RENDERER]);
-      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);
   }
 
-
   /**
    * Returns the root wrapper
    *
@@ -227,8 +228,7 @@ class ShallowWrapper {
       throw new Error('ShallowWrapper::update() can only be called on the root');
     }
     this.single('update', () => {
-      this[NODE] = getRootNode(this[RENDERER].getNode());
-      this[NODES] = [this[NODE]];
+      privateSetNodes(this, getRootNode(this[RENDERER].getNode()));
     });
     return this;
   }