From 14a6c6c227e02527b614c4ef9d110c324edeb7b1 Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@gmail.com>
Date: Mon, 11 Jun 2018 20:03:51 +0100
Subject: [PATCH] Allow multiple root children in test renderer traversal API
 (#13017)

---
 src/ReactTestRenderer.js                      | 91 ++++++++++++-------
 .../ReactTestRendererTraversal-test.js        | 44 +++++++++
 2 files changed, 101 insertions(+), 34 deletions(-)

diff --git a/src/ReactTestRenderer.js b/src/ReactTestRenderer.js
index 3fff46d..711f8f6 100644
--- a/src/ReactTestRenderer.js
+++ b/src/ReactTestRenderer.js
@@ -200,8 +200,45 @@ const validWrapperTypes = new Set([
   ClassComponent,
   HostComponent,
   ForwardRef,
+  // Normally skipped, but used when there's more than one root child.
+  HostRoot,
 ]);
 
+function getChildren(parent: Fiber) {
+  const children = [];
+  const startingNode = parent;
+  let node: Fiber = startingNode;
+  if (node.child === null) {
+    return children;
+  }
+  node.child.return = node;
+  node = node.child;
+  outer: while (true) {
+    let descend = false;
+    if (validWrapperTypes.has(node.tag)) {
+      children.push(wrapFiber(node));
+    } else if (node.tag === HostText) {
+      children.push('' + node.memoizedProps);
+    } else {
+      descend = true;
+    }
+    if (descend && node.child !== null) {
+      node.child.return = node;
+      node = node.child;
+      continue;
+    }
+    while (node.sibling === null) {
+      if (node.return === startingNode) {
+        break outer;
+      }
+      node = (node.return: any);
+    }
+    (node.sibling: any).return = node.return;
+    node = (node.sibling: any);
+  }
+  return children;
+}
+
 class ReactTestInstance {
   _fiber: Fiber;
 
@@ -246,6 +283,13 @@ class ReactTestInstance {
     let parent = this._fiber.return;
     while (parent !== null) {
       if (validWrapperTypes.has(parent.tag)) {
+        if (parent.tag === HostRoot) {
+          // Special case: we only "materialize" instances for roots
+          // if they have more than a single child. So we'll check that now.
+          if (getChildren(parent).length < 2) {
+            return null;
+          }
+        }
         return wrapFiber(parent);
       }
       parent = parent.return;
@@ -254,38 +298,7 @@ class ReactTestInstance {
   }
 
   get children(): Array<ReactTestInstance | string> {
-    const children = [];
-    const startingNode = this._currentFiber();
-    let node: Fiber = startingNode;
-    if (node.child === null) {
-      return children;
-    }
-    node.child.return = node;
-    node = node.child;
-    outer: while (true) {
-      let descend = false;
-      if (validWrapperTypes.has(node.tag)) {
-        children.push(wrapFiber(node));
-      } else if (node.tag === HostText) {
-        children.push('' + node.memoizedProps);
-      } else {
-        descend = true;
-      }
-      if (descend && node.child !== null) {
-        node.child.return = node;
-        node = node.child;
-        continue;
-      }
-      while (node.sibling === null) {
-        if (node.return === startingNode) {
-          break outer;
-        }
-        node = (node.return: any);
-      }
-      (node.sibling: any).return = node.return;
-      node = (node.sibling: any);
-    }
-    return children;
+    return getChildren(this._currentFiber());
   }
 
   // Custom search functions
@@ -469,10 +482,20 @@ const ReactTestRendererFiber = {
         configurable: true,
         enumerable: true,
         get: function() {
-          if (root === null || root.current.child === null) {
+          if (root === null) {
+            throw new Error("Can't access .root on unmounted test renderer");
+          }
+          const children = getChildren(root.current);
+          if (children.length === 0) {
             throw new Error("Can't access .root on unmounted test renderer");
+          } else if (children.length === 1) {
+            // Normally, we skip the root and just give you the child.
+            return children[0];
+          } else {
+            // However, we give you the root if there's more than one root child.
+            // We could make this the behavior for all cases but it would be a breaking change.
+            return wrapFiber(root.current);
           }
-          return wrapFiber(root.current.child);
         },
       }: Object),
     );
diff --git a/src/__tests__/ReactTestRendererTraversal-test.js b/src/__tests__/ReactTestRendererTraversal-test.js
index 9fb91e1..3ea7188 100644
--- a/src/__tests__/ReactTestRendererTraversal-test.js
+++ b/src/__tests__/ReactTestRendererTraversal-test.js
@@ -199,4 +199,48 @@ describe('ReactTestRendererTraversal', () => {
     expect(nestedViews[1].parent).toBe(expectedParent);
     expect(nestedViews[2].parent).toBe(expectedParent);
   });
+
+  it('can have special nodes as roots', () => {
+    const FR = React.forwardRef(props => <section {...props} />);
+    expect(
+      ReactTestRenderer.create(
+        <FR>
+          <div />
+          <div />
+        </FR>,
+      ).root.findAllByType('div').length,
+    ).toBe(2);
+    expect(
+      ReactTestRenderer.create(
+        <React.Fragment>
+          <div />
+          <div />
+        </React.Fragment>,
+      ).root.findAllByType('div').length,
+    ).toBe(2);
+    expect(
+      ReactTestRenderer.create(
+        <React.Fragment key="foo">
+          <div />
+          <div />
+        </React.Fragment>,
+      ).root.findAllByType('div').length,
+    ).toBe(2);
+    expect(
+      ReactTestRenderer.create(
+        <React.StrictMode>
+          <div />
+          <div />
+        </React.StrictMode>,
+      ).root.findAllByType('div').length,
+    ).toBe(2);
+    expect(
+      ReactTestRenderer.create(
+        <Context.Provider>
+          <div />
+          <div />
+        </Context.Provider>,
+      ).root.findAllByType('div').length,
+    ).toBe(2);
+  });
 });