Add traversal for Fiber test renderer

Not clear the path to shipping this but this gives us a migration path internally that we need right now (replaces https://fburl.com/udq9ksvk).

Author: sophiebits

src/renderers/testing/ReactTestRendererFeatureFlags.js

@@ -0,0 +1,17 @@
+/**
+ * Copyright 2013-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree. An additional grant
+ * of patent rights can be found in the PATENTS file in the same directory.
+ *
+ * @providesModule ReactTestRendererFeatureFlags
+ * @flow
+ */
+
+const ReactTestRendererFeatureFlags = {
+  enableTraversal: false,
+};
+
+module.exports = ReactTestRendererFeatureFlags;

src/renderers/testing/ReactTestRendererFiberEntry.js

@@ -14,11 +14,14 @@
 'use strict';
 
 var ReactFiberReconciler = require('ReactFiberReconciler');
+var ReactFiberTreeReflection = require('ReactFiberTreeReflection');
 var ReactGenericBatching = require('ReactGenericBatching');
+var ReactTestRendererFeatureFlags = require('ReactTestRendererFeatureFlags');
 var emptyObject = require('fbjs/lib/emptyObject');
 var ReactTypeOfWork = require('ReactTypeOfWork');
 var invariant = require('fbjs/lib/invariant');
 var {
+  Fragment,
   FunctionalComponent,
   ClassComponent,
   HostComponent,
@@ -63,6 +66,19 @@ type TextInstance = {|
 
 const UPDATE_SIGNAL = {};
 
+function getPublicInstance(inst: Instance | TextInstance): * {
+  switch (inst.tag) {
+    case 'INSTANCE':
+      const createNodeMock = inst.rootContainerInstance.createNodeMock;
+      return createNodeMock({
+        type: inst.type,
+        props: inst.props,
+      });
+    default:
+      return inst;
+  }
+}
+
 function appendChild(
   parentInstance: Instance | Container,
   child: Instance | TextInstance,
@@ -225,18 +241,7 @@ var TestRenderer = ReactFiberReconciler({
 
   useSyncScheduling: true,
 
-  getPublicInstance(inst: Instance | TextInstance): * {
-    switch (inst.tag) {
-      case 'INSTANCE':
-        const createNodeMock = inst.rootContainerInstance.createNodeMock;
-        return createNodeMock({
-          type: inst.type,
-          props: inst.props,
-        });
-      default:
-        return inst;
-    }
-  },
+  getPublicInstance,
 });
 
 var defaultTestOptions = {
@@ -325,6 +330,193 @@ function toTree(node: ?Fiber) {
   }
 }
 
+const fiberToWrapper = new WeakMap();
+function wrapFiber(fiber: Fiber): ReactTestInstance {
+  let wrapper = fiberToWrapper.get(fiber);
+  if (wrapper === undefined && fiber.alternate !== null) {
+    wrapper = fiberToWrapper.get(fiber.alternate);
+  }
+  if (wrapper === undefined) {
+    wrapper = new ReactTestInstance(fiber);
+    fiberToWrapper.set(fiber, wrapper);
+  }
+  return wrapper;
+}
+
+const validWrapperTypes = new Set([
+  FunctionalComponent,
+  ClassComponent,
+  HostComponent,
+]);
+
+class ReactTestInstance {
+  _fiber: Fiber;
+
+  _currentFiber(): Fiber {
+    const mounted = ReactFiberTreeReflection.isFiberMounted(this._fiber);
+    invariant(mounted, 'Can\'t read properties from unmounted component');
+    return ReactFiberTreeReflection.findCurrentFiberUsingSlowPath(this._fiber);
+  }
+
+  constructor(fiber: Fiber) {
+    invariant(
+      validWrapperTypes.has(fiber.tag),
+      'Unexpected object passed to ReactTestInstance constructor (tag: %s). ' +
+      'This is probably a bug in React.',
+      fiber.tag
+    );
+    this._fiber = fiber;
+  }
+
+  get instance() {
+    if (this._fiber.tag === HostComponent) {
+      return getPublicInstance(this._fiber.stateNode);
+    } else {
+      return this._fiber.stateNode;
+    }
+  }
+
+  get type() {
+    return this._fiber.type;
+  }
+
+  get props(): Object {
+    return this._currentFiber().memoizedProps;
+  }
+
+  get parent(): Object {
+    const parent = this._fiber.return;
+    return parent.tag === HostRoot ? null : wrapFiber(parent);
+  }
+
+  get children(): Array<ReactTestInstance | string> {
+    const children = [];
+    const startingNode = this._currentFiber();
+    if (startingNode.child === null) {
+      return children;
+    }
+    let node = startingNode;
+    node.child.return = node;
+    node = node.child;
+    while (true) {
+      let descend = false;
+      switch (node.tag) {
+        case FunctionalComponent:
+        case ClassComponent:
+        case HostComponent:
+          children.push(wrapFiber(node));
+          break;
+        case HostText:
+          children.push('' + node.memoizedProps);
+          break;
+        case Fragment:
+          descend = true;
+          break;
+        default:
+          invariant(
+            false,
+            'Unsupported component type %s in test renderer. ' +
+            'This is probably a bug in React.',
+            node.tag
+          );
+      }
+      if (descend && node.child !== null) {
+        node.child.return = node;
+        node = node.child;
+        continue;
+      }
+      while (node.sibling === null) {
+        if (node.return === startingNode) {
+          return children;
+        }
+        node = node.return;
+      }
+      node.sibling.return = node.return;
+      node = node.sibling;
+    }
+  }
+
+  // Custom search functions
+  find(predicate: Predicate): ReactTestInstance {
+    return expectOne(
+      this.findAll(predicate, {deep: false}),
+      `matching custom predicate: ${predicate.toString()}`
+    );
+  }
+
+  findByType(type: any): ReactTestInstance {
+    return expectOne(
+      this.findAllByType(type, {deep: false}),
+      `with node type: "${type.displayName || type.name}"`
+    );
+  }
+
+  findByProps(props: Object): ReactTestInstance {
+    return expectOne(
+      this.findAllByProps(props, {deep: false}),
+      `with props: ${JSON.stringify(props)}`
+    );
+  }
+
+  findAll(predicate: Predicate, options: ?FindOptions = null): Array<ReactTestInstance> {
+    return findAll(this, predicate, options);
+  }
+
+  findAllByType(type: any, options: ?FindOptions = null): Array<ReactTestInstance> {
+    return findAll(this, node => node.type === type, options);
+  }
+
+  findAllByProps(props: Object, options: ?FindOptions = null): Array<ReactTestInstance> {
+    return findAll(this, node => node.props && propsMatch(node.props, props), options);
+  }
+}
+
+function findAll(
+  root: ReactTestInstance,
+  predicate: Predicate,
+  options: ?FindOptions,
+): Array<ReactTestInstance> {
+  const deep = options ? options.deep : true;
+  const results = [];
+
+  if (predicate(root)) {
+    results.push(root);
+    if (!deep) {
+      return results;
+    }
+  }
+
+  for (const child of root.children) {
+    if (typeof child === 'string') {
+      continue;
+    }
+    results.push(...findAll(child, predicate, options));
+  }
+
+  return results;
+}
+
+function expectOne(all: Array<ReactTestInstance>, message: string): ReactTestInstance {
+  if (all.length === 1) {
+    return all[0];
+  }
+
+  const prefix = all.length === 0
+    ? 'No instances found '
+    : `Expected 1 but found ${all.length} instances `;
+
+  throw new Error(prefix + message);
+}
+
+function propsMatch(props: Object, filter: Object): boolean {
+  for (const key in filter) {
+    if (props[key] !== filter[key]) {
+      return false;
+    }
+  }
+  return true;
+}
+
 var ReactTestRendererFiber = {
   create(element: ReactElement<any>, options: TestRendererOptions) {
     var createNodeMock = defaultTestOptions.createNodeMock;
@@ -341,6 +533,17 @@ var ReactTestRendererFiber = {
     TestRenderer.updateContainer(element, root, null, null);
 
     return {
+      get root() {
+        if (!ReactTestRendererFeatureFlags.enableTraversal) {
+          throw new Error(
+            'Test renderer traversal is experimental and not enabled'
+          );
+        }
+        if (root === null) {
+          throw new Error('Can\'t access .root on unmounted test renderer');
+        }
+        return wrapFiber(root.current.child);
+      },
       toJSON() {
         if (root == null || root.current == null || container == null) {
           return null;