Add unstable APIs for async rendering to test renderer

packages/react-native-renderer/src/ReactNativeFrameScheduling.js

@@ -23,6 +23,7 @@ let frameDeadline: number = 0;
 
 const frameDeadlineObject: Deadline = {
   timeRemaining: () => frameDeadline - now(),
+  didTimeout: false,
 };
 
 function setTimeoutCallback() {

packages/react-noop-renderer/src/ReactNoop.js

@@ -308,6 +308,10 @@ function* flushUnitsOfWork(n: number): Generator<Array<mixed>, void, void> {
         didStop = true;
         return 0;
       },
+      // React's scheduler has its own way of keeping track of expired
+      // work and doesn't read this, so don't bother setting it to the
+      // correct value.
+      didTimeout: false,
     });
 
     if (yieldedValues !== null) {

packages/react-reconciler/src/ReactFiberReconciler.js

@@ -37,6 +37,7 @@ if (__DEV__) {
 
 export type Deadline = {
   timeRemaining: () => number,
+  didTimeout: boolean,
 };
 
 type OpaqueHandle = Fiber;

packages/react-test-renderer/src/ReactTestRenderer.js

@@ -9,6 +9,7 @@
 
 import type {Fiber} from 'react-reconciler/src/ReactFiber';
 import type {FiberRoot} from 'react-reconciler/src/ReactFiberRoot';
+import type {Deadline} from 'react-reconciler/src/ReactFiberReconciler';
 
 import ReactFiberReconciler from 'react-reconciler';
 import {batchedUpdates} from 'events/ReactGenericBatching';
@@ -31,6 +32,7 @@ import invariant from 'fbjs/lib/invariant';
 
 type TestRendererOptions = {
   createNodeMock: (element: React$Element<any>) => any,
+  unstable_isAsync: boolean,
 };
 
 type ReactTestRendererJSON = {|
@@ -116,6 +118,11 @@ function removeChild(
   parentInstance.children.splice(index, 1);
 }
 
+// Current virtual time
+let currentTime: number = 0;
+let scheduledCallback: ((deadline: Deadline) => mixed) | null = null;
+let yieldedValues: Array<mixed> | null = null;
+
 const TestRenderer = ReactFiberReconciler({
   getRootHostContext() {
     return emptyObject;
@@ -200,19 +207,22 @@ const TestRenderer = ReactFiberReconciler({
     };
   },
 
-  scheduleDeferredCallback(fn: Function): number {
-    return setTimeout(fn, 0, {timeRemaining: Infinity});
+  scheduleDeferredCallback(
+    callback: (deadline: Deadline) => mixed,
+    options?: {timeout: number},
+  ): number {
+    scheduledCallback = callback;
+    return 0;
   },
 
   cancelDeferredCallback(timeoutID: number): void {
-    clearTimeout(timeoutID);
+    scheduledCallback = null;
   },
 
   getPublicInstance,
 
   now(): number {
-    // Test renderer does not use expiration
-    return 0;
+    return currentTime;
   },
 
   mutation: {
@@ -603,8 +613,14 @@ function propsMatch(props: Object, filter: Object): boolean {
 const ReactTestRendererFiber = {
   create(element: React$Element<any>, options: TestRendererOptions) {
     let createNodeMock = defaultTestOptions.createNodeMock;
-    if (options && typeof options.createNodeMock === 'function') {
-      createNodeMock = options.createNodeMock;
+    let isAsync = false;
+    if (typeof options === 'object' && options !== null) {
+      if (typeof options.createNodeMock === 'function') {
+        createNodeMock = options.createNodeMock;
+      }
+      if (options.unstable_isAsync === true) {
+        isAsync = true;
+      }
     }
     let container = {
       children: [],
@@ -613,7 +629,7 @@ const ReactTestRendererFiber = {
     };
     let root: FiberRoot | null = TestRenderer.createContainer(
       container,
-      false,
+      isAsync,
       false,
     );
     invariant(root != null, 'something went wrong');
@@ -654,6 +670,66 @@ const ReactTestRendererFiber = {
         container = null;
         root = null;
       },
+      unstable_flushAll(): Array<mixed> {
+        yieldedValues = null;
+        while (scheduledCallback !== null) {
+          const cb = scheduledCallback;
+          scheduledCallback = null;
+          cb({
+            timeRemaining() {
+              // Keep rendering until there's no more work
+              return 999;
+            },
+            // React's scheduler has its own way of keeping track of expired
+            // work and doesn't read this, so don't bother setting it to the
+            // correct value.
+            didTimeout: false,
+          });
+        }
+        if (yieldedValues === null) {
+          // Always return an array.
+          return [];
+        }
+        return yieldedValues;
+      },
+      unstable_flushThrough(expectedValues: Array<mixed>): Array<mixed> {
+        let didStop = false;
+        yieldedValues = null;
+        while (scheduledCallback !== null && !didStop) {
+          const cb = scheduledCallback;
+          scheduledCallback = null;
+          cb({
+            timeRemaining() {
+              if (
+                yieldedValues !== null &&
+                yieldedValues.length >= expectedValues.length
+              ) {
+                // We at least as many values as expected. Stop rendering.
+                didStop = true;
+                return 0;
+              }
+              // Keep rendering.
+              return 999;
+            },
+            // React's scheduler has its own way of keeping track of expired
+            // work and doesn't read this, so don't bother setting it to the
+            // correct value.
+            didTimeout: false,
+          });
+        }
+        if (yieldedValues === null) {
+          // Always return an array.
+          return [];
+        }
+        return yieldedValues;
+      },
+      unstable_yield(value: mixed): void {
+        if (yieldedValues === null) {
+          yieldedValues = [value];
+        } else {
+          yieldedValues.push(value);
+        }
+      },
       getInstance() {
         if (root == null || root.current == null) {
           return null;

packages/react-test-renderer/src/__tests__/ReactTestRendererAsync-test.js

@@ -0,0 +1,97 @@
+/**
+ * Copyright (c) 2013-present, Facebook, Inc.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails react-core
+ * @jest-environment node
+ */
+
+'use strict';
+
+const React = require('react');
+const ReactTestRenderer = require('react-test-renderer');
+
+describe('ReactTestRendererAsync', () => {
+  it('flushAll flushes all work', () => {
+    function Foo(props) {
+      return props.children;
+    }
+    const renderer = ReactTestRenderer.create(<Foo>Hi</Foo>, {
+      unstable_isAsync: true,
+    });
+
+    // Before flushing, nothing has mounted.
+    expect(renderer.toJSON()).toEqual(null);
+
+    // Flush initial mount.
+    renderer.unstable_flushAll();
+    expect(renderer.toJSON()).toEqual('Hi');
+
+    // Update
+    renderer.update(<Foo>Bye</Foo>);
+    // Not yet updated.
+    expect(renderer.toJSON()).toEqual('Hi');
+    // Flush update.
+    renderer.unstable_flushAll();
+    expect(renderer.toJSON()).toEqual('Bye');
+  });
+
+  it('flushAll returns array of yielded values', () => {
+    function Child(props) {
+      renderer.unstable_yield(props.children);
+      return props.children;
+    }
+    function Parent(props) {
+      return (
+        <React.Fragment>
+          <Child>{'A:' + props.step}</Child>
+          <Child>{'B:' + props.step}</Child>
+          <Child>{'C:' + props.step}</Child>
+        </React.Fragment>
+      );
+    }
+    const renderer = ReactTestRenderer.create(<Parent step={1} />, {
+      unstable_isAsync: true,
+    });
+
+    expect(renderer.unstable_flushAll()).toEqual(['A:1', 'B:1', 'C:1']);
+    expect(renderer.toJSON()).toEqual(['A:1', 'B:1', 'C:1']);
+
+    renderer.update(<Parent step={2} />);
+    expect(renderer.unstable_flushAll()).toEqual(['A:2', 'B:2', 'C:2']);
+    expect(renderer.toJSON()).toEqual(['A:2', 'B:2', 'C:2']);
+  });
+
+  it('flushThrough flushes until the expected values is yielded', () => {
+    function Child(props) {
+      renderer.unstable_yield(props.children);
+      return props.children;
+    }
+    function Parent(props) {
+      return (
+        <React.Fragment>
+          <Child>{'A:' + props.step}</Child>
+          <Child>{'B:' + props.step}</Child>
+          <Child>{'C:' + props.step}</Child>
+        </React.Fragment>
+      );
+    }
+    const renderer = ReactTestRenderer.create(<Parent step={1} />, {
+      unstable_isAsync: true,
+    });
+
+    // Flush the first two siblings
+    expect(renderer.unstable_flushThrough(['A:1', 'B:1'])).toEqual([
+      'A:1',
+      'B:1',
+    ]);
+    // Did not commit yet.
+    expect(renderer.toJSON()).toEqual(null);
+
+    // Flush the remaining work
+    expect(renderer.unstable_flushAll()).toEqual(['C:1']);
+    expect(renderer.toJSON()).toEqual(['A:1', 'B:1', 'C:1']);
+  });
+});

packages/shared/ReactDOMFrameScheduling.js

@@ -63,6 +63,7 @@ if (!ExecutionEnvironment.canUseDOM) {
         timeRemaining() {
           return Infinity;
         },
+        didTimeout: false,
       });
     });
   };

packages/shared/forks/ReactFeatureFlags.test-renderer.js

@@ -0,0 +1,36 @@
+/**
+ * Copyright (c) 2013-present, Facebook, Inc.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import invariant from 'fbjs/lib/invariant';
+
+import typeof * as FeatureFlagsType from 'shared/ReactFeatureFlags';
+import typeof * as PersistentFeatureFlagsType from './ReactFeatureFlags.persistent';
+
+export const debugRenderPhaseSideEffects = false;
+export const debugRenderPhaseSideEffectsForStrictMode = false;
+export const enableCreateRoot = false;
+export const enableUserTimingAPI = __DEV__;
+export const enableGetDerivedStateFromCatch = false;
+export const warnAboutDeprecatedLifecycles = false;
+export const replayFailedUnitOfWorkWithInvokeGuardedCallback = false;
+export const enableMutatingReconciler = true;
+export const enableNoopReconciler = false;
+export const enablePersistentReconciler = false;
+export const alwaysUseRequestIdleCallbackPolyfill = false;
+
+// Only used in www builds.
+export function addUserTimingListener() {
+  invariant(false, 'Not implemented.');
+}
+
+// Flow magic to verify the exports of this file match the original version.
+// eslint-disable-next-line no-unused-vars
+type Check<_X, Y: _X, X: Y = _X> = null;
+// eslint-disable-next-line no-unused-expressions
+(null: Check<PersistentFeatureFlagsType, FeatureFlagsType>);

scripts/rollup/forks.js

@@ -38,6 +38,8 @@ const forks = Object.freeze({
         return 'shared/forks/ReactFeatureFlags.native-fabric.js';
       case 'react-reconciler/persistent':
         return 'shared/forks/ReactFeatureFlags.persistent.js';
+      case 'react-test-renderer':
+        return 'shared/forks/ReactFeatureFlags.test-renderer.js';
       default:
         switch (bundleType) {
           case FB_DEV: