[Fiber] Fix TestUtils edge cases

scripts/fiber/tests-failing.txt

@@ -593,9 +593,7 @@ src/shared/utils/__tests__/traverseAllChildren-test.js
 * should warn for using maps as children with owner info
 
 src/test/__tests__/ReactTestUtils-test.js
-* traverses children in the correct order
 * should support injected wrapper components as DOM components
 * should change the value of an input field
 * should change the value of an input field in a component
-* can scry with stateless components involved
 * should set the type of the event

scripts/fiber/tests-passing.txt

@@ -1260,8 +1260,10 @@ src/test/__tests__/ReactTestUtils-test.js
 * can scryRenderedDOMComponentsWithClass with TextComponent
 * can scryRenderedDOMComponentsWithClass with className contains \n
 * can scryRenderedDOMComponentsWithClass with multiple classes
+* traverses children in the correct order
 * should throw when attempting to use ReactTestUtils.Simulate with shallow rendering
 * should not warn when simulating events with extra properties
+* can scry with stateless components involved
 
 src/test/__tests__/reactComponentExpect-test.js
 * should match composite components

src/test/ReactTestUtils.js

@@ -32,7 +32,9 @@ var invariant = require('invariant');
 var topLevelTypes = EventConstants.topLevelTypes;
 var {
   ClassComponent,
+  FunctionalComponent,
   HostComponent,
+  HostContainer,
   HostText,
 } = ReactTypeOfWork;
 
@@ -80,6 +82,7 @@ function findAllInRenderedFiberTreeInternal(fiber, test) {
   }
   if (
     fiber.tag !== ClassComponent &&
+    fiber.tag !== FunctionalComponent &&
     fiber.tag !== HostComponent &&
     fiber.tag !== HostText
   ) {
@@ -209,7 +212,15 @@ var ReactTestUtils = {
     );
     var internalInstance = ReactInstanceMap.get(inst);
     if (internalInstance && typeof internalInstance.tag === 'number') {
-      return findAllInRenderedFiberTreeInternal(internalInstance, test);
+      var fiber = internalInstance;
+      var root = fiber;
+      while (root.return) {
+        root = root.return;
+      }
+      var isRootCurrent = root.tag === HostContainer && root.stateNode.current === root;
+      // Make sure we're introspecting the current tree
+      var current = isRootCurrent ? fiber : fiber.alternate;
+      return findAllInRenderedFiberTreeInternal(current, test);
     } else {
       return findAllInRenderedStackTreeInternal(internalInstance, test);
     }