Fix fragment handling in toTree() (#12154)

Author: gaearon

src/ReactTestRenderer.js

@@ -19,6 +19,7 @@ import {
   FunctionalComponent,
   ClassComponent,
   HostComponent,
+  HostPortal,
   HostText,
   HostRoot,
 } from 'shared/ReactTypeOfWork';
@@ -286,60 +287,86 @@ function toJSON(inst: Instance | TextInstance): ReactTestRendererNode {
   }
 }
 
-function nodeAndSiblingsTrees(nodeWithSibling: ?Fiber) {
+function childrenToTree(node) {
+  if (!node) {
+    return null;
+  }
+  const children = nodeAndSiblingsArray(node);
+  if (children.length === 0) {
+    return null;
+  } else if (children.length === 1) {
+    return toTree(children[0]);
+  }
+  return flatten(children.map(toTree));
+}
+
+function nodeAndSiblingsArray(nodeWithSibling) {
   const array = [];
   let node = nodeWithSibling;
   while (node != null) {
     array.push(node);
     node = node.sibling;
   }
-  const trees = array.map(toTree);
-  return trees.length ? trees : null;
+  return array;
 }
 
-function hasSiblings(node: ?Fiber) {
-  return node && node.sibling;
+function flatten(arr) {
+  const result = [];
+  const stack = [{i: 0, array: arr}];
+  while (stack.length) {
+    const n = stack.pop();
+    while (n.i < n.array.length) {
+      const el = n.array[n.i];
+      n.i += 1;
+      if (Array.isArray(el)) {
+        stack.push(n);
+        stack.push({i: 0, array: el});
+        break;
+      }
+      result.push(el);
+    }
+  }
+  return result;
 }
 
 function toTree(node: ?Fiber) {
   if (node == null) {
     return null;
   }
   switch (node.tag) {
-    case HostRoot: // 3
-      return toTree(node.child);
+    case HostRoot:
+      return childrenToTree(node.child);
+    case HostPortal:
+      return childrenToTree(node.child);
     case ClassComponent:
       return {
         nodeType: 'component',
         type: node.type,
         props: {...node.memoizedProps},
         instance: node.stateNode,
-        rendered: hasSiblings(node.child)
-          ? nodeAndSiblingsTrees(node.child)
-          : toTree(node.child),
+        rendered: childrenToTree(node.child),
       };
-    case FunctionalComponent: // 1
+    case FunctionalComponent:
       return {
         nodeType: 'component',
         type: node.type,
         props: {...node.memoizedProps},
         instance: null,
-        rendered: hasSiblings(node.child)
-          ? nodeAndSiblingsTrees(node.child)
-          : toTree(node.child),
+        rendered: childrenToTree(node.child),
       };
-    case HostComponent: // 5
+    case HostComponent: {
       return {
         nodeType: 'host',
         type: node.type,
         props: {...node.memoizedProps},
         instance: null, // TODO: use createNodeMock here somehow?
-        rendered: nodeAndSiblingsTrees(node.child),
+        rendered: flatten(nodeAndSiblingsArray(node.child).map(toTree)),
       };
-    case HostText: // 6
+    }
+    case HostText:
       return node.stateNode.text;
-    case Fragment: // 10
-      return toTree(node.child);
+    case Fragment:
+      return childrenToTree(node.child);
     default:
       invariant(
         false,

src/__tests__/ReactTestRenderer-test.js

@@ -615,7 +615,7 @@ describe('ReactTestRenderer', () => {
     );
   });
 
-  it('toTree() handles complicated tree of fragments', () => {
+  it('toTree() handles complicated tree of arrays', () => {
     class Foo extends React.Component {
       render() {
         return this.props.children;
@@ -693,6 +693,58 @@ describe('ReactTestRenderer', () => {
     );
   });
 
+  it('toTree() handles complicated tree of fragments', () => {
+    const renderer = ReactTestRenderer.create(
+      <React.Fragment>
+        <React.Fragment>
+          <div>One</div>
+          <div>Two</div>
+          <React.Fragment>
+            <div>Three</div>
+          </React.Fragment>
+        </React.Fragment>
+        <div>Four</div>
+      </React.Fragment>,
+    );
+
+    const tree = renderer.toTree();
+
+    cleanNodeOrArray(tree);
+
+    expect(prettyFormat(tree)).toEqual(
+      prettyFormat([
+        {
+          type: 'div',
+          nodeType: 'host',
+          props: {},
+          instance: null,
+          rendered: ['One'],
+        },
+        {
+          type: 'div',
+          nodeType: 'host',
+          props: {},
+          instance: null,
+          rendered: ['Two'],
+        },
+        {
+          type: 'div',
+          nodeType: 'host',
+          props: {},
+          instance: null,
+          rendered: ['Three'],
+        },
+        {
+          type: 'div',
+          nodeType: 'host',
+          props: {},
+          instance: null,
+          rendered: ['Four'],
+        },
+      ]),
+    );
+  });
+
   it('root instance and createNodeMock ref return the same value', () => {
     const createNodeMock = ref => ({node: ref});
     let refInst = null;