From fc90eb636876d54d99ace2773dd4923f3e848106 Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Tue, 28 Mar 2023 00:03:57 -0400
Subject: [PATCH] Codemod more tests to waitFor pattern (#26494)

---
 .../src/__tests__/ReactUpdates-test.js        |  19 ++--
 .../__tests__/ReactIncrementalUpdates-test.js | 105 ++++++++++++------
 2 files changed, 78 insertions(+), 46 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactUpdates-test.js b/packages/react-dom/src/__tests__/ReactUpdates-test.js
index 3c900a1987b2d..67abeb6dd0ad6 100644
--- a/packages/react-dom/src/__tests__/ReactUpdates-test.js
+++ b/packages/react-dom/src/__tests__/ReactUpdates-test.js
@@ -1627,7 +1627,6 @@ describe('ReactUpdates', () => {
         const [step, setStep] = React.useState(0);
         React.useEffect(() => {
           setStep(x => x + 1);
-          Scheduler.log(step);
         });
         return step;
       }
@@ -1642,23 +1641,19 @@ describe('ReactUpdates', () => {
       console.error = (e, s) => {
         error = e;
         stack = s;
+        Scheduler.log('stop');
       };
       try {
         const container = document.createElement('div');
-        expect(() => {
-          const root = ReactDOMClient.createRoot(container);
-          root.render(<App />);
-          while (error === null) {
-            Scheduler.unstable_flushNumberOfYields(1);
-            Scheduler.unstable_clearLog();
-          }
-          expect(stack).toContain(' NonTerminating');
-          // rethrow error to prevent going into an infinite loop when act() exits
-          throw error;
-        }).toThrow('Maximum update depth exceeded.');
+        const root = ReactDOMClient.createRoot(container);
+        root.render(<App />);
+        await waitFor(['stop']);
       } finally {
         console.error = originalConsoleError;
       }
+
+      expect(error).toContain('Maximum update depth exceeded');
+      expect(stack).toContain('at NonTerminating');
     });
 
     it('can have nested updates if they do not cross the limit', async () => {
diff --git a/packages/react-reconciler/src/__tests__/ReactIncrementalUpdates-test.js b/packages/react-reconciler/src/__tests__/ReactIncrementalUpdates-test.js
index 5fd12fee1f8bb..a85cf89c68ca4 100644
--- a/packages/react-reconciler/src/__tests__/ReactIncrementalUpdates-test.js
+++ b/packages/react-reconciler/src/__tests__/ReactIncrementalUpdates-test.js
@@ -36,13 +36,9 @@ describe('ReactIncrementalUpdates', () => {
     assertLog = InternalTestUtils.assertLog;
   });
 
-  function flushNextRenderIfExpired() {
-    // This will start rendering the next level of work. If the work hasn't
-    // expired yet, React will exit without doing anything. If it has expired,
-    // it will schedule a sync task.
-    Scheduler.unstable_flushExpired();
-    // Flush the sync task.
-    ReactNoop.flushSync();
+  function Text({text}) {
+    Scheduler.log(text);
+    return text;
   }
 
   it('applies updates in order of priority', async () => {
@@ -528,35 +524,38 @@ describe('ReactIncrementalUpdates', () => {
       setCount = _setCount;
       Scheduler.log('Render: ' + count);
       useLayoutEffect(() => {
-        setCount(prevCount => prevCount + 1);
+        setCount(1);
         Scheduler.log('Commit: ' + count);
       }, []);
-      return null;
+      return <Text text="Child" />;
     }
 
     await act(async () => {
       React.startTransition(() => {
         ReactNoop.render(<App />);
       });
-      flushNextRenderIfExpired();
       assertLog([]);
-      await waitForAll(['Render: 0', 'Commit: 0', 'Render: 1']);
+      await waitForAll([
+        'Render: 0',
+        'Child',
+        'Commit: 0',
+        'Render: 1',
+        'Child',
+      ]);
 
       Scheduler.unstable_advanceTime(10000);
       React.startTransition(() => {
         setCount(2);
       });
-      flushNextRenderIfExpired();
-      assertLog([]);
+      // The transition should not have expired, so we should be able to
+      // partially render it.
+      await waitFor(['Render: 2']);
+      // Now do the rest
+      await waitForAll(['Child']);
     });
   });
 
-  it('regression: does not expire soon due to previous flushSync', () => {
-    function Text({text}) {
-      Scheduler.log(text);
-      return text;
-    }
-
+  it('regression: does not expire soon due to previous flushSync', async () => {
     ReactNoop.flushSync(() => {
       ReactNoop.render(<Text text="A" />);
     });
@@ -565,32 +564,70 @@ describe('ReactIncrementalUpdates', () => {
     Scheduler.unstable_advanceTime(10000);
 
     React.startTransition(() => {
-      ReactNoop.render(<Text text="B" />);
+      ReactNoop.render(
+        <>
+          <Text text="A" />
+          <Text text="B" />
+          <Text text="C" />
+          <Text text="D" />
+        </>,
+      );
+    });
+    // The transition should not have expired, so we should be able to
+    // partially render it.
+    await waitFor(['A']);
+
+    // FIXME:  We should be able to partially render B, too, but currently it
+    // expires. This is an existing bug that I discovered, which will be fixed
+    // in a PR that I'm currently working on.
+    //
+    // Correct behavior:
+    //   await waitFor(['B']);
+    //   await waitForAll(['C', 'D']);
+    //
+    // Current behavior:
+    await waitFor(['B'], {
+      additionalLogsAfterAttemptingToYield: ['C', 'D'],
     });
-    flushNextRenderIfExpired();
-    assertLog([]);
   });
 
-  it('regression: does not expire soon due to previous expired work', () => {
-    function Text({text}) {
-      Scheduler.log(text);
-      return text;
-    }
-
+  it('regression: does not expire soon due to previous expired work', async () => {
     React.startTransition(() => {
-      ReactNoop.render(<Text text="A" />);
+      ReactNoop.render(
+        <>
+          <Text text="A" />
+          <Text text="B" />
+          <Text text="C" />
+          <Text text="D" />
+        </>,
+      );
     });
+    await waitFor(['A']);
+
+    // This will expire the rest of the update
     Scheduler.unstable_advanceTime(10000);
-    flushNextRenderIfExpired();
-    assertLog(['A']);
+    await waitFor(['B'], {
+      additionalLogsAfterAttemptingToYield: ['C', 'D'],
+    });
 
     Scheduler.unstable_advanceTime(10000);
 
+    // Now do another transition. This one should not expire.
     React.startTransition(() => {
-      ReactNoop.render(<Text text="B" />);
+      ReactNoop.render(
+        <>
+          <Text text="A" />
+          <Text text="B" />
+          <Text text="C" />
+          <Text text="D" />
+        </>,
+      );
     });
-    flushNextRenderIfExpired();
-    assertLog([]);
+    // The transition should not have expired, so we should be able to
+    // partially render it.
+    await waitFor(['A']);
+    await waitFor(['B']);
+    await waitForAll(['C', 'D']);
   });
 
   it('when rebasing, does not exclude updates that were already committed, regardless of priority', async () => {