From 3412c7da902cf273e49a18de98c973f1acc9ff5b Mon Sep 17 00:00:00 2001
From: Ben Alpert <balpert@fb.com>
Date: Thu, 3 Aug 2017 19:04:34 -0700
Subject: [PATCH] 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).
---
 .flowconfig                                   |   1 +
 .../testing/ReactTestRendererFeatureFlags.js  |  17 ++
 .../testing/ReactTestRendererFiberEntry.js    | 263 +++++++++++++++++-
 .../ReactTestRendererTraversal-test.js        | 157 +++++++++++
 4 files changed, 425 insertions(+), 13 deletions(-)
 create mode 100644 src/renderers/testing/ReactTestRendererFeatureFlags.js
 create mode 100644 src/renderers/testing/__tests__/ReactTestRendererTraversal-test.js

diff --git a/.flowconfig b/.flowconfig
index f11ad7d01cded..e6985a2146c12 100644
--- a/.flowconfig
+++ b/.flowconfig
@@ -24,6 +24,7 @@ module.system=haste
 
 esproposal.class_static_fields=enable
 esproposal.class_instance_fields=enable
+unsafe.enable_getters_and_setters=true
 
 munge_underscores=false
 
diff --git a/src/renderers/testing/ReactTestRendererFeatureFlags.js b/src/renderers/testing/ReactTestRendererFeatureFlags.js
new file mode 100644
index 0000000000000..6b405899941a9
--- /dev/null
+++ b/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;
diff --git a/src/renderers/testing/ReactTestRendererFiberEntry.js b/src/renderers/testing/ReactTestRendererFiberEntry.js
index d290d704812b5..b08ddb2ec10bb 100644
--- a/src/renderers/testing/ReactTestRendererFiberEntry.js
+++ b/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,
@@ -61,8 +64,29 @@ type TextInstance = {|
   tag: 'TEXT',
 |};
 
+type FindOptions = $Shape<{
+  // performs a "greedy" search: if a matching node is found, will continue
+  // to search within the matching node's children. (default: true)
+  deep: boolean,
+}>;
+
+export type Predicate = (node: ReactTestInstance) => ?boolean;
+
 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 +249,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 +338,219 @@ 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 {
+    // Throws if this component has been unmounted.
+    const fiber = ReactFiberTreeReflection.findCurrentFiberUsingSlowPath(
+      this._fiber,
+    );
+    invariant(
+      fiber !== null,
+      "Can't read from currently-mounting component. This error is likely " +
+        'caused by a bug in React. Please file an issue.',
+    );
+    return 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(): ?ReactTestInstance {
+    const parent = this._fiber.return;
+    return parent === null || parent.tag === HostRoot
+      ? null
+      : wrapFiber(parent);
+  }
+
+  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;
+      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) {
+          break outer;
+        }
+        node = (node.return: any);
+      }
+      (node.sibling: any).return = node.return;
+      node = (node.sibling: any);
+    }
+    return children;
+  }
+
+  // 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;
@@ -336,11 +562,22 @@ var ReactTestRendererFiber = {
       createNodeMock,
       tag: 'CONTAINER',
     };
-    var root: ?FiberRoot = TestRenderer.createContainer(container);
+    var root: FiberRoot | null = TestRenderer.createContainer(container);
     invariant(root != null, 'something went wrong');
     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 || root.current.child === 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;
diff --git a/src/renderers/testing/__tests__/ReactTestRendererTraversal-test.js b/src/renderers/testing/__tests__/ReactTestRendererTraversal-test.js
new file mode 100644
index 0000000000000..1f43adcb29f9b
--- /dev/null
+++ b/src/renderers/testing/__tests__/ReactTestRendererTraversal-test.js
@@ -0,0 +1,157 @@
+/**
+ * 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.
+ *
+ * @emails react-core
+ */
+
+'use strict';
+
+const React = require('react');
+let ReactTestRenderer;
+
+const RCTView = 'RCTView';
+const View = props => <RCTView {...props} />;
+
+describe('ReactTestRendererTraversal', () => {
+  beforeEach(() => {
+    jest.resetModules();
+    const ReactTestRendererFeatureFlags = require('ReactTestRendererFeatureFlags');
+    ReactTestRendererFeatureFlags.enableTraversal = true;
+    ReactTestRenderer = require('react-test-renderer');
+  });
+
+  class Example extends React.Component {
+    render() {
+      return (
+        <View>
+          <View foo="foo">
+            <View bar="bar" />
+            <View bar="bar" baz="baz" itself="itself" />
+            <View />
+            <ExampleSpread bar="bar" />
+            <ExampleFn bar="bar" bing="bing" />
+            <ExampleNull bar="bar" />
+            <ExampleNull null="null">
+              <View void="void" />
+              <View void="void" />
+            </ExampleNull>
+          </View>
+        </View>
+      );
+    }
+  }
+  class ExampleSpread extends React.Component {
+    render = () => <View {...this.props} />;
+  }
+  const ExampleFn = props => <View baz="baz" />;
+  const ExampleNull = props => null;
+
+  it('initializes', () => {
+    const render = ReactTestRenderer.create(<Example />);
+    const hasFooProp = node => node.props.hasOwnProperty('foo');
+
+    // assert .props, .type and .parent attributes
+    const foo = render.root.find(hasFooProp);
+    expect(foo.props.children).toHaveLength(7);
+    expect(foo.type).toBe(View);
+    expect(render.root.parent).toBe(null);
+    expect(foo.children[0].parent).toBe(foo);
+  });
+
+  it('searches via .find() / .findAll()', () => {
+    const render = ReactTestRenderer.create(<Example />);
+    const hasFooProp = node => node.props.hasOwnProperty('foo');
+    const hasBarProp = node => node.props.hasOwnProperty('bar');
+    const hasBazProp = node => node.props.hasOwnProperty('baz');
+    const hasBingProp = node => node.props.hasOwnProperty('bing');
+    const hasNullProp = node => node.props.hasOwnProperty('null');
+    const hasVoidProp = node => node.props.hasOwnProperty('void');
+    const hasItselfProp = node => node.props.hasOwnProperty('itself');
+
+    expect(() => render.root.find(hasFooProp)).not.toThrow(); // 1 match
+    expect(() => render.root.find(hasBarProp)).toThrow(); // >1 matches
+    expect(() => render.root.find(hasBazProp)).toThrow(); // >1 matches
+    expect(() => render.root.find(hasBingProp)).not.toThrow(); // 1 match
+    expect(() => render.root.find(hasNullProp)).not.toThrow(); // 1 match
+    expect(() => render.root.find(hasVoidProp)).toThrow(); // 0 matches
+
+    // same assertion as .find(), but confirm length
+    expect(render.root.findAll(hasFooProp, {deep: false})).toHaveLength(1);
+    expect(render.root.findAll(hasBarProp, {deep: false})).toHaveLength(5);
+    expect(render.root.findAll(hasBazProp, {deep: false})).toHaveLength(2);
+    expect(render.root.findAll(hasBingProp, {deep: false})).toHaveLength(1);
+    expect(render.root.findAll(hasNullProp, {deep: false})).toHaveLength(1);
+    expect(render.root.findAll(hasVoidProp, {deep: false})).toHaveLength(0);
+
+    // note: with {deep: true}, .findAll() will continue to
+    //       search children, even after finding a match
+    expect(render.root.findAll(hasFooProp)).toHaveLength(2);
+    expect(render.root.findAll(hasBarProp)).toHaveLength(9);
+    expect(render.root.findAll(hasBazProp)).toHaveLength(4);
+    expect(render.root.findAll(hasBingProp)).toHaveLength(1); // no spread
+    expect(render.root.findAll(hasNullProp)).toHaveLength(1); // no spread
+    expect(render.root.findAll(hasVoidProp)).toHaveLength(0);
+
+    const bing = render.root.find(hasBingProp);
+    expect(bing.find(hasBarProp)).toBe(bing);
+    expect(bing.find(hasBingProp)).toBe(bing);
+    expect(bing.findAll(hasBazProp, {deep: false})).toHaveLength(1);
+    expect(bing.findAll(hasBazProp)).toHaveLength(2);
+
+    const foo = render.root.find(hasFooProp);
+    expect(foo.findAll(hasFooProp, {deep: false})).toHaveLength(1);
+    expect(foo.findAll(hasFooProp)).toHaveLength(2);
+
+    const itself = foo.find(hasItselfProp);
+    expect(itself.find(hasBarProp)).toBe(itself);
+    expect(itself.find(hasBazProp)).toBe(itself);
+    expect(itself.findAll(hasBazProp, {deep: false})).toHaveLength(1);
+    expect(itself.findAll(hasBazProp)).toHaveLength(2);
+  });
+
+  it('searches via .findByType() / .findAllByType()', () => {
+    const render = ReactTestRenderer.create(<Example />);
+
+    expect(() => render.root.findByType(ExampleFn)).not.toThrow(); // 1 match
+    expect(() => render.root.findByType(View)).not.toThrow(); // 1 match
+    // note: there are clearly multiple <View /> in general, but there
+    //       is only one being rendered at root node level
+    expect(() => render.root.findByType(ExampleNull)).toThrow(); // 2 matches
+
+    expect(render.root.findAllByType(ExampleFn)).toHaveLength(1);
+    expect(render.root.findAllByType(View, {deep: false})).toHaveLength(1);
+    expect(render.root.findAllByType(View)).toHaveLength(7);
+    expect(render.root.findAllByType(ExampleNull)).toHaveLength(2);
+
+    const nulls = render.root.findAllByType(ExampleNull);
+    expect(nulls[0].findAllByType(View)).toHaveLength(0);
+    expect(nulls[1].findAllByType(View)).toHaveLength(0);
+
+    const fn = render.root.findAllByType(ExampleFn);
+    expect(fn[0].findAllByType(View)).toHaveLength(1);
+  });
+
+  it('searches via .findByProps() / .findAllByProps()', () => {
+    const render = ReactTestRenderer.create(<Example />);
+    const foo = 'foo';
+    const bar = 'bar';
+    const baz = 'baz';
+
+    expect(() => render.root.findByProps({foo})).not.toThrow(); // 1 match
+    expect(() => render.root.findByProps({bar})).toThrow(); // >1 matches
+    expect(() => render.root.findByProps({baz})).toThrow(); // >1 matches
+
+    expect(render.root.findAllByProps({foo}, {deep: false})).toHaveLength(1);
+    expect(render.root.findAllByProps({bar}, {deep: false})).toHaveLength(5);
+    expect(render.root.findAllByProps({baz}, {deep: false})).toHaveLength(2);
+
+    expect(render.root.findAllByProps({foo})).toHaveLength(2);
+    expect(render.root.findAllByProps({bar})).toHaveLength(9);
+    expect(render.root.findAllByProps({baz})).toHaveLength(4);
+  });
+});