Convert a few more tests to waitFor test helpers (#26509)

Continuing my journey to migrate all the Scheduler flush* methods to
async versions of the same helpers.

packages/internal-test-utils/ReactInternalTestUtils.js

@@ -12,25 +12,25 @@ import enqueueTask from './enqueueTask';
 
 export {act} from './internalAct';
 
-function assertYieldsWereCleared(Scheduler) {
-  const actualYields = Scheduler.unstable_clearLog();
+function assertYieldsWereCleared(caller) {
+  const actualYields = SchedulerMock.unstable_clearLog();
   if (actualYields.length !== 0) {
     const error = Error(
       'The event log is not empty. Call assertLog(...) first.',
     );
-    Error.captureStackTrace(error, assertYieldsWereCleared);
+    Error.captureStackTrace(error, caller);
     throw error;
   }
 }
 
-async function waitForMicrotasks() {
+export async function waitForMicrotasks() {
   return new Promise(resolve => {
     enqueueTask(() => resolve());
   });
 }
 
 export async function waitFor(expectedLog, options) {
-  assertYieldsWereCleared(SchedulerMock);
+  assertYieldsWereCleared(waitFor);
 
   // Create the error object before doing any async work, to get a better
   // stack trace.
@@ -79,7 +79,7 @@ ${diff(expectedLog, actualLog)}
 }
 
 export async function waitForAll(expectedLog) {
-  assertYieldsWereCleared(SchedulerMock);
+  assertYieldsWereCleared(waitForAll);
 
   // Create the error object before doing any async work, to get a better
   // stack trace.
@@ -110,7 +110,7 @@ ${diff(expectedLog, actualLog)}
 }
 
 export async function waitForThrow(expectedError: mixed): mixed {
-  assertYieldsWereCleared(SchedulerMock);
+  assertYieldsWereCleared(waitForThrow);
 
   // Create the error object before doing any async work, to get a better
   // stack trace.
@@ -160,7 +160,7 @@ ${diff(expectedError, x)}
 // avoid using it in tests. It's really only for testing a particular
 // implementation detail (update starvation prevention).
 export async function unstable_waitForExpired(expectedLog): mixed {
-  assertYieldsWereCleared(SchedulerMock);
+  assertYieldsWereCleared(unstable_waitForExpired);
 
   // Create the error object before doing any async work, to get a better
   // stack trace.
@@ -189,7 +189,7 @@ ${diff(expectedLog, actualLog)}
 // now because that's how untable_flushUntilNextPaint already worked, but maybe
 // we should split these use cases into separate APIs.
 export async function waitForPaint(expectedLog) {
-  assertYieldsWereCleared(SchedulerMock);
+  assertYieldsWereCleared(waitForPaint);
 
   // Create the error object before doing any async work, to get a better
   // stack trace.
@@ -219,7 +219,7 @@ ${diff(expectedLog, actualLog)}
 }
 
 export async function waitForDiscrete(expectedLog) {
-  assertYieldsWereCleared(SchedulerMock);
+  assertYieldsWereCleared(waitForDiscrete);
 
   // Create the error object before doing any async work, to get a better
   // stack trace.

packages/react-dom/src/__tests__/ReactMultiChildText-test.js

@@ -10,11 +10,12 @@
 'use strict';
 
 const React = require('react');
-const ReactDOM = require('react-dom');
+const ReactDOMClient = require('react-dom/client');
 const ReactTestUtils = require('react-dom/test-utils');
+const act = require('internal-test-utils').act;
 
 // Helpers
-const testAllPermutations = function (testCases) {
+const testAllPermutations = async function (testCases) {
   for (let i = 0; i < testCases.length; i += 2) {
     const renderWithChildren = testCases[i];
     const expectedResultAfterRender = testCases[i + 1];
@@ -24,10 +25,11 @@ const testAllPermutations = function (testCases) {
       const expectedResultAfterUpdate = testCases[j + 1];
 
       const container = document.createElement('div');
-      ReactDOM.render(<div>{renderWithChildren}</div>, container);
+      const root = ReactDOMClient.createRoot(container);
+      await act(() => root.render(<div>{renderWithChildren}</div>));
       expectChildren(container, expectedResultAfterRender);
 
-      ReactDOM.render(<div>{updateWithChildren}</div>, container);
+      await act(() => root.render(<div>{updateWithChildren}</div>));
       expectChildren(container, expectedResultAfterUpdate);
     }
   }
@@ -75,10 +77,12 @@ const expectChildren = function (container, children) {
  * faster to render and update.
  */
 describe('ReactMultiChildText', () => {
-  it('should correctly handle all possible children for render and update', () => {
-    expect(() => {
+  jest.setTimeout(20000);
+
+  it('should correctly handle all possible children for render and update', async () => {
+    await expect(async () => {
       // prettier-ignore
-      testAllPermutations([
+      await testAllPermutations([
         // basic values
         undefined, [],
         null, [],

packages/react-reconciler/src/__tests__/ReactHooks-test.internal.js

@@ -1703,17 +1703,14 @@ describe('ReactHooks', () => {
       return null;
     }
 
-    await act(() => {
-      ReactTestRenderer.unstable_batchedUpdates(() => {
-        ReactTestRenderer.create(
-          <>
-            <A />
-            <B />
-          </>,
-        );
-        expect(() => Scheduler.unstable_flushAll()).toThrow('Hello');
-      });
-    });
+    expect(() => {
+      ReactTestRenderer.create(
+        <>
+          <A />
+          <B />
+        </>,
+      );
+    }).toThrow('Hello');
 
     if (__DEV__) {
       expect(console.error).toHaveBeenCalledTimes(2);

packages/react-reconciler/src/__tests__/ReactIsomorphicAct-test.js

@@ -16,6 +16,7 @@ let use;
 let Suspense;
 let DiscreteEventPriority;
 let startTransition;
+let waitForMicrotasks;
 
 describe('isomorphic act()', () => {
   beforeEach(() => {
@@ -28,6 +29,8 @@ describe('isomorphic act()', () => {
     use = React.use;
     Suspense = React.Suspense;
     startTransition = React.startTransition;
+
+    waitForMicrotasks = require('internal-test-utils').waitForMicrotasks;
   });
 
   beforeEach(() => {
@@ -51,7 +54,7 @@ describe('isomorphic act()', () => {
     // Nothing has rendered yet
     expect(root).toMatchRenderedOutput(null);
     // Flush the microtasks by awaiting
-    await null;
+    await waitForMicrotasks();
     expect(root).toMatchRenderedOutput('A');
 
     // Now do the same thing but wrap the update with `act`. No
@@ -229,9 +232,7 @@ describe('isomorphic act()', () => {
     //
     // The exact number of microtasks is an implementation detail; just needs
     // to happen when the microtask queue is flushed.
-    await null;
-    await null;
-    await null;
+    await waitForMicrotasks();
 
     expect(console.error).toHaveBeenCalledTimes(1);
     expect(console.error.mock.calls[0][0]).toContain(
@@ -282,9 +283,7 @@ describe('isomorphic act()', () => {
     //
     // The exact number of microtasks is an implementation detail; just needs
     // to happen when the microtask queue is flushed.
-    await null;
-    await null;
-    await null;
+    await waitForMicrotasks();
 
     expect(console.error).toHaveBeenCalledTimes(0);
 

packages/react-reconciler/src/__tests__/ReactSchedulerIntegration-test.js

@@ -102,6 +102,7 @@ describe('ReactSchedulerIntegration', () => {
     const root = ReactNoop.createRoot();
     root.render('Initial');
     await waitForAll([]);
+    expect(root).toMatchRenderedOutput('Initial');
 
     scheduleCallback(NormalPriority, () => Scheduler.log('A'));
     scheduleCallback(NormalPriority, () => Scheduler.log('B'));
@@ -112,16 +113,22 @@ describe('ReactSchedulerIntegration', () => {
       root.render('Update');
     });
 
-    // Advance time just to be sure the next tasks have lower priority
-    Scheduler.unstable_advanceTime(2000);
+    // Perform just a little bit of work. By now, the React task will have
+    // already been scheduled, behind A, B, and C.
+    await waitFor(['A']);
 
+    // Schedule some additional tasks. These won't fire until after the React
+    // update has finished.
     scheduleCallback(NormalPriority, () => Scheduler.log('D'));
     scheduleCallback(NormalPriority, () => Scheduler.log('E'));
 
     // Flush everything up to the next paint. Should yield after the
     // React commit.
-    Scheduler.unstable_flushUntilNextPaint();
-    assertLog(['A', 'B', 'C']);
+    await waitForPaint(['B', 'C']);
+    expect(root).toMatchRenderedOutput('Update');
+
+    // Now flush the rest of the work.
+    await waitForAll(['D', 'E']);
   });
 
   // @gate www
@@ -213,6 +220,7 @@ describe(
       waitForPaint = InternalTestUtils.waitForPaint;
       assertLog = InternalTestUtils.assertLog;
       waitFor = InternalTestUtils.waitFor;
+      act = InternalTestUtils.act;
     });
 
     afterEach(() => {
@@ -293,8 +301,7 @@ describe(
         // Start logging whenever shouldYield is called
         logDuringShouldYield = true;
         // Let's call it once to confirm the mock actually works
-        Scheduler.unstable_shouldYield();
-        assertLog(['shouldYield']);
+        await waitFor(['shouldYield']);
 
         // Expire the task
         Scheduler.unstable_advanceTime(10000);
@@ -307,11 +314,9 @@ describe(
         startTransition(() => {
           ReactNoop.render(<App />);
         });
-
         // Because the render expired, React should finish the tree without
         // consulting `shouldYield` again
-        Scheduler.unstable_flushNumberOfYields(1);
-        assertLog(['B', 'C']);
+        await waitFor(['B', 'C']);
       });
     });
   },