From 928541dc9ee1df266d621541922f71bb6e4b770f Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@gmail.com>
Date: Sat, 29 Oct 2016 22:02:13 +0100
Subject: [PATCH] [Fiber] Fix TestUtils.findAllInRenderedTree (#8147)

* [Fiber] Fix TestUtils.findAllInRenderedTree

* Add a comment about coroutines
---
 src/test/ReactTestUtils.js | 46 ++++++++++++++++++++++++++++++++++----
 1 file changed, 42 insertions(+), 4 deletions(-)

diff --git a/src/test/ReactTestUtils.js b/src/test/ReactTestUtils.js
index a5c3b31d23..0c2e567ba7 100644
--- a/src/test/ReactTestUtils.js
+++ b/src/test/ReactTestUtils.js
@@ -20,6 +20,7 @@ var ReactDOM = require('ReactDOM');
 var ReactDOMComponentTree = require('ReactDOMComponentTree');
 var ReactBrowserEventEmitter = require('ReactBrowserEventEmitter');
 var ReactInstanceMap = require('ReactInstanceMap');
+var ReactTypeOfWork = require('ReactTypeOfWork');
 var ReactUpdates = require('ReactUpdates');
 var SyntheticEvent = require('SyntheticEvent');
 var ReactShallowRenderer = require('ReactShallowRenderer');
@@ -28,6 +29,11 @@ var findDOMNode = require('findDOMNode');
 var invariant = require('invariant');
 
 var topLevelTypes = EventConstants.topLevelTypes;
+var {
+  ClassComponent,
+  HostComponent,
+  HostText,
+} = ReactTypeOfWork;
 
 function Event(suffix) {}
 
@@ -35,7 +41,7 @@ function Event(suffix) {}
  * @class ReactTestUtils
  */
 
-function findAllInRenderedTreeInternal(inst, test) {
+function findAllInRenderedStackTreeInternal(inst, test) {
   if (!inst || !inst.getPublicInstance) {
     return [];
   }
@@ -50,7 +56,7 @@ function findAllInRenderedTreeInternal(inst, test) {
         continue;
       }
       ret = ret.concat(
-        findAllInRenderedTreeInternal(
+        findAllInRenderedStackTreeInternal(
           renderedChildren[key],
           test
         )
@@ -61,12 +67,39 @@ function findAllInRenderedTreeInternal(inst, test) {
     typeof currentElement.type === 'function'
   ) {
     ret = ret.concat(
-      findAllInRenderedTreeInternal(inst._renderedComponent, test)
+      findAllInRenderedStackTreeInternal(inst._renderedComponent, test)
     );
   }
   return ret;
 }
 
+function findAllInRenderedFiberTreeInternal(fiber, test) {
+  if (!fiber) {
+    return [];
+  }
+  if (
+    fiber.tag !== ClassComponent &&
+    fiber.tag !== HostComponent &&
+    fiber.tag !== HostText
+  ) {
+    return [];
+  }
+  var publicInst = fiber.stateNode;
+  var ret = publicInst && test(publicInst) ? [publicInst] : [];
+  var child = fiber.child;
+  while (child) {
+    ret = ret.concat(
+      findAllInRenderedFiberTreeInternal(
+        child,
+        test
+      )
+    );
+    child = child.sibling;
+  }
+  // TODO: visit stateNode for coroutines
+  return ret;
+}
+
 /**
  * Utilities for making it easy to test React components.
  *
@@ -170,7 +203,12 @@ var ReactTestUtils = {
       ReactTestUtils.isCompositeComponent(inst),
       'findAllInRenderedTree(...): instance must be a composite component'
     );
-    return findAllInRenderedTreeInternal(ReactInstanceMap.get(inst), test);
+    var internalInstance = ReactInstanceMap.get(inst);
+    if (internalInstance && typeof internalInstance.tag === 'number') {
+      return findAllInRenderedFiberTreeInternal(internalInstance, test);
+    } else {
+      return findAllInRenderedStackTreeInternal(internalInstance, test);
+    }
   },
 
   /**