Fix failing tests by using act

Co-authored-by: Andrew Clark <hi@andrewclark.io>

Author: tyao1

packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js

@@ -449,10 +449,9 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(deleted.length).toBe(0);
 
     // Performing an update should force it to delete the boundary
-    root.render(<App value={true} />);
-
-    Scheduler.unstable_flushAll();
-    jest.runAllTimers();
+    await act(async () => {
+      root.render(<App value={true} />);
+    });
 
     expect(hydrated.length).toBe(1);
     expect(deleted.length).toBe(1);
@@ -945,13 +944,12 @@ describe('ReactDOMServerPartialHydration', () => {
     root.render(<App text="Hi" className="hi" />);
 
     // At the same time, resolving the promise so that rendering can complete.
-    suspend = false;
-    resolve();
-    await promise;
-
     // This should first complete the hydration and then flush the update onto the hydrated state.
-    Scheduler.unstable_flushAll();
-    jest.runAllTimers();
+    await act(async () => {
+      suspend = false;
+      resolve();
+      await promise;
+    });
 
     // The new span should be the same since we should have successfully hydrated
     // before changing it.
@@ -1093,9 +1091,9 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(ref.current).toBe(null);
 
     // Render an update, but leave it still suspended.
-    root.render(<App text="Hi" className="hi" />);
-    Scheduler.unstable_flushAll();
-    jest.runAllTimers();
+    await act(async () => {
+      root.render(<App text="Hi" className="hi" />);
+    });
 
     // Flushing now should delete the existing content and show the fallback.
 
@@ -1104,12 +1102,11 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(container.textContent).toBe('Loading...');
 
     // Unsuspending shows the content.
-    suspend = false;
-    resolve();
-    await promise;
-
-    Scheduler.unstable_flushAll();    
-    jest.runAllTimers();
+    await act(async () => {
+      suspend = false;
+      resolve();
+      await promise;
+    });
 
     const span = container.getElementsByTagName('span')[0];
     expect(span.textContent).toBe('Hi');
@@ -1174,23 +1171,21 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(ref.current).toBe(span);
 
     // Render an update, but leave it still suspended.
-    root.render(<App text="Hi" className="hi" />);
-
     // Flushing now should delete the existing content and show the fallback.
-    Scheduler.unstable_flushAll();
-    jest.runAllTimers();
+    await act(async () => {
+      root.render(<App text="Hi" className="hi" />);
+    });
 
     expect(container.getElementsByTagName('span').length).toBe(1);
     expect(ref.current).toBe(span);
     expect(container.textContent).toBe('');
 
     // Unsuspending shows the content.
-    suspend = false;
-    resolve();
-    await promise;
-
-    Scheduler.unstable_flushAll();
-    jest.runAllTimers();
+    await act(async () => {
+      suspend = false;
+      resolve();
+      await promise;
+    });
 
     expect(span.textContent).toBe('Hi');
     expect(span.className).toBe('hi');
@@ -1252,20 +1247,21 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(ref.current).toBe(null);
 
     // Render an update, but leave it still suspended.
-    root.render(<App text="Hi" className="hi" />);
-
     // Flushing now should delete the existing content and show the fallback.
-    Scheduler.unstable_flushAll();
-    jest.runAllTimers();
+    await act(async () => {
+      root.render(<App text="Hi" className="hi" />);
+    });
 
     expect(container.getElementsByTagName('span').length).toBe(0);
     expect(ref.current).toBe(null);
     expect(container.textContent).toBe('Loading...');
 
     // Unsuspending shows the content.
-    suspend = false;
-    resolve();
-    await promise;
+    await act(async () => {
+      suspend = false;
+      resolve();
+      await promise;
+    });
 
     Scheduler.unstable_flushAll();
     jest.runAllTimers();
@@ -1490,13 +1486,12 @@ describe('ReactDOMServerPartialHydration', () => {
     );
 
     // At the same time, resolving the promise so that rendering can complete.
-    suspend = false;
-    resolve();
-    await promise;
-
     // This should first complete the hydration and then flush the update onto the hydrated state.
-    Scheduler.unstable_flushAll();
-    jest.runAllTimers();
+    await act(async () => {
+      suspend = false;
+      resolve();
+      await promise;
+    });
 
     // Since this should have been hydrated, this should still be the same span.
     const newSpan = container.getElementsByTagName('span')[0];
@@ -1569,27 +1564,25 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(ref.current).toBe(null);
 
     // Render an update, but leave it still suspended.
-    root.render(
-      <Context.Provider value={{text: 'Hi', className: 'hi'}}>
-        <App />
-      </Context.Provider>,
-    );
-
     // Flushing now should delete the existing content and show the fallback.
-    Scheduler.unstable_flushAll();
-    jest.runAllTimers();
+    await act(async () => {
+      root.render(
+        <Context.Provider value={{text: 'Hi', className: 'hi'}}>
+          <App />
+        </Context.Provider>,
+      );
+    });
 
     expect(container.getElementsByTagName('span').length).toBe(0);
     expect(ref.current).toBe(null);
     expect(container.textContent).toBe('Loading...');
 
     // Unsuspending shows the content.
-    suspend = false;
-    resolve();
-    await promise;
-
-    Scheduler.unstable_flushAll();
-    jest.runAllTimers();
+    await act(async () => {
+      suspend = false;
+      resolve();
+      await promise;
+    });
 
     const span = container.getElementsByTagName('span')[0];
     expect(span.textContent).toBe('Hi');
@@ -2320,16 +2313,15 @@ describe('ReactDOMServerPartialHydration', () => {
 
     // Render an update, which will be higher or the same priority as pinging the hydration.
     // The new update doesn't suspend.
-    root.render(
-      <ClassName.Provider value={'hi'}>
-        <App text="Hi" />
-      </ClassName.Provider>,
-    );
-
     // Since we're still suspended on the original data, we can't hydrate.
     // This will force all expiration times to flush.
-    Scheduler.unstable_flushAll();
-    jest.runAllTimers();
+    await act(async () => {
+      root.render(
+        <ClassName.Provider value={'hi'}>
+          <App text="Hi" />
+        </ClassName.Provider>,
+      );
+    });
 
     // This will now be a new span because we weren't able to hydrate before
     const newSpan = container.getElementsByTagName('span')[0];

packages/react-dom/src/__tests__/ReactDOMServerSelectiveHydration-test.internal.js

@@ -1786,7 +1786,7 @@ describe('ReactDOMServerSelectiveHydration', () => {
     document.body.removeChild(container);
   });
 
-  it('can force hydration in response to sync update', () => {
+  it('can force hydration in response to sync update', async () => {
     function Child({text}) {
       Scheduler.unstable_yieldValue(`Child ${text}`);
       return <span ref={ref => (spanRef = ref)}>{text}</span>;
@@ -1812,15 +1812,17 @@ describe('ReactDOMServerSelectiveHydration', () => {
     const root = ReactDOMClient.hydrateRoot(container, <App text="A" />);
     expect(Scheduler).toFlushUntilNextPaint(['App A']);
 
-    ReactDOM.flushSync(() => {
-      root.render(<App text="B" />);
+    await act(async () => {
+      ReactDOM.flushSync(() => {
+        root.render(<App text="B" />);
+      });
     });
     expect(Scheduler).toHaveYielded(['App B', 'Child A', 'App B', 'Child B']);
     expect(initialSpan).toBe(spanRef);
   });
 
   // @gate experimental || www
-  it('can force hydration in response to continuous update', () => {
+  it('can force hydration in response to continuous update', async () => {
     function Child({text}) {
       Scheduler.unstable_yieldValue(`Child ${text}`);
       return <span ref={ref => (spanRef = ref)}>{text}</span>;
@@ -1846,14 +1848,17 @@ describe('ReactDOMServerSelectiveHydration', () => {
     const root = ReactDOMClient.hydrateRoot(container, <App text="A" />);
     expect(Scheduler).toFlushUntilNextPaint(['App A']);
 
-    TODO_scheduleContinuousSchedulerTask(() => {
-      root.render(<App text="B" />);
+    await act(async () => {
+      TODO_scheduleContinuousSchedulerTask(() => {
+        root.render(<App text="B" />);
+      });
     });
-    expect(Scheduler).toFlushAndYield(['App B', 'Child A', 'App B', 'Child B']);
+
+    expect(Scheduler).toHaveYielded(['App B', 'Child A', 'App B', 'Child B']);
     expect(initialSpan).toBe(spanRef);
   });
 
-  it('can force hydration in response to default update', () => {
+  it('can force hydration in response to default update', async () => {
     function Child({text}) {
       Scheduler.unstable_yieldValue(`Child ${text}`);
       return <span ref={ref => (spanRef = ref)}>{text}</span>;
@@ -1878,11 +1883,10 @@ describe('ReactDOMServerSelectiveHydration', () => {
     const initialSpan = container.getElementsByTagName('span')[0];
     const root = ReactDOMClient.hydrateRoot(container, <App text="A" />);
     expect(Scheduler).toFlushUntilNextPaint(['App A']);
-
-    ReactDOM.unstable_batchedUpdates(() => {
+    await act(async () => {
       root.render(<App text="B" />);
     });
-    expect(Scheduler).toFlushAndYield(['App B', 'Child A', 'App B', 'Child B']);
+    expect(Scheduler).toHaveYielded(['App B', 'Child A', 'App B', 'Child B']);
     expect(initialSpan).toBe(spanRef);
   });
 

packages/react-reconciler/src/ReactFiberLane.js

@@ -766,14 +766,11 @@ export function getBumpedLaneForHydration(
   const renderLane = getHighestPriorityLane(renderLanes);
 
   let lane;
-  /*
-  TODO: 
   if (enableUnifiedSyncLane) {
     if ((renderLane & SyncUpdateLanes) !== NoLane) {
       lane = SyncHydrationLane;
     }
   }
-  */
   if (!lane) {
     switch (renderLane) {
       case SyncLane: