Repository: facebook/react
File: packages/react-reconciler/src/__tests__/ReactSuspense-test.internal.js (JavaScript)

Code before:
        });
      }

      function App({shouldSuspend, step}) {
        return (
          <>
            <Text text={`A${step}`} />
            <Suspense fallback={<Text text="Loading..." />}>
              {shouldSuspend ? <AsyncText text="Async" ms={2000} /> : null}
            </Suspense>
            <Text text={`B${step}`} />
            <Text text={`C${step}`} />
          </>
        );
      }

      await act(() => {
        root.render(null);
      });

      root.render(<App shouldSuspend={false} step={0} />);
      await waitForAll(['A0', 'B0', 'C0']);
      expect(container.textContent).toEqual('A0B0C0');

      // This update will suspend.
      await act(() => {
        root.render(<App shouldSuspend={true} step={1} />);
      });

      // Do a bit of work
      await waitFor(['A1']);

      // Schedule another update. This will have lower priority because it's
      // a transition.
      React.startTransition(async () => {
        await act(() => {
          root.render(<App shouldSuspend={false} step={2} />);
        });
      });

      // Interrupt to trigger a restart.
      await interrupt();

      await waitFor([
        // Should have restarted the first update, because of the interruption
        'A1',
        'Suspend! [Async]',
        'Loading...',
        'B1',
      ]);


Review comment (suggestion):
`act` will flush everything, but the `waitFor` only flushes up to a certain log. I would expect flushing everything here to break something, since the point of this test is to partially flush, then schedule an update.

Code after:

      function App({shouldSuspend, step}) {
        return (
          <>
            <Text text={`A${step}`} />
            <Suspense fallback={<Text text="Loading..." />}>
              {shouldSuspend ? <AsyncText text="Async" ms={2000} /> : null}
            </Suspense>
            <Text text={`B${step}`} />
            <Text text={`C${step}`} />
          </>
        );
      }

      root.render(<App shouldSuspend={false} step={0} />);
      await waitForAll(['A0', 'B0', 'C0']);
      expect(container.textContent).toEqual('A0B0C0');

      // This update will suspend.
      root.render(<App shouldSuspend={true} step={1} />);

      // Do a bit of work
      await waitFor(['A1']);

      // Schedule another update. This will have lower priority because it's
      // a transition.
      React.startTransition(() => {
        root.render(<App shouldSuspend={false} step={2} />);
      });

      // Interrupt to trigger a restart.
      interrupt();

      await waitFor([
        // Should have restarted the first update, because of the interruption
        'A1',
        'Suspend! [Async]',
        'Loading...',
        'B1',
      ]);

      // Should not have committed loading state
      expect(container.textContent).toEqual('A0B0C0');

      // After suspending, should abort the first update and switch to the
      // second update. So, C1 should not appear in the log.
      // TODO: This should work even if React does not yield to the main
      // thread. Should use same mechanism as selective hydration to interrupt
      // the render before the end of the current slice of work.
      await waitForAll(['A2', 'B2', 'C2']);