From da03de3884915d390bf69858ed56454a1e96fb51 Mon Sep 17 00:00:00 2001
From: Matt Brophy <matt@brophy.org>
Date: Tue, 2 May 2023 18:02:22 -0400
Subject: [PATCH 1/5] Wrap internal router state updates with
 React.startTransition

---
 .changeset/start-transition-support.md        |   6 +
 .../concurrent-mode-navigations-test.tsx      | 347 ++++++++++++++++++
 packages/react-router-dom/index.tsx           |  38 +-
 packages/react-router/lib/components.tsx      |  22 +-
 4 files changed, 404 insertions(+), 9 deletions(-)
 create mode 100644 .changeset/start-transition-support.md
 create mode 100644 packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx

diff --git a/.changeset/start-transition-support.md b/.changeset/start-transition-support.md
new file mode 100644
index 0000000000..e4f3d54a48
--- /dev/null
+++ b/.changeset/start-transition-support.md
@@ -0,0 +1,6 @@
+---
+"react-router": minor
+"react-router-dom": minor
+---
+
+Wrap internal router state updates with `React.startTransition` if it exists
diff --git a/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx b/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx
new file mode 100644
index 0000000000..ae0b1008d6
--- /dev/null
+++ b/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx
@@ -0,0 +1,347 @@
+import * as React from "react";
+import {
+  MemoryRouter,
+  Routes,
+  Route,
+  useNavigate,
+  BrowserRouter,
+  HashRouter,
+  createMemoryRouter,
+  createRoutesFromElements,
+  RouterProvider,
+} from "react-router-dom";
+import {
+  act,
+  fireEvent,
+  prettyDOM,
+  render,
+  screen,
+  waitFor,
+} from "@testing-library/react";
+import { JSDOM } from "jsdom";
+
+describe("Handles concurrent mode features during navigations", () => {
+  function getComponents(withBoundary: boolean) {
+    function Home() {
+      let navigate = useNavigate();
+      return (
+        <>
+          <h1>Home</h1>
+          <button onClick={() => navigate("/about")}>/about</button>
+        </>
+      );
+    }
+
+    let resolved = false;
+    let dfd = createDeferred();
+    function resolve() {
+      resolved = true;
+      return dfd.resolve();
+    }
+
+    function AboutWithoutBoundary() {
+      if (!resolved) {
+        throw dfd.promise;
+      }
+      return <h1>About</h1>;
+    }
+
+    function AboutWithBoundary() {
+      if (!resolved) {
+        throw dfd.promise;
+      }
+      return (
+        <React.Suspense fallback={<p>Loading...</p>}>
+          <h1>About</h1>
+        </React.Suspense>
+      );
+    }
+
+    return {
+      Home,
+      About: withBoundary ? AboutWithBoundary : AboutWithoutBoundary,
+      resolve,
+    };
+  }
+
+  describe("when the destination route suspends with a boundary", () => {
+    async function assertNavigation(
+      container: HTMLElement,
+      resolve: () => void
+    ) {
+      expect(getHtml(container)).toMatchInlineSnapshot(`
+        "<div>
+          <h1>
+            Home
+          </h1>
+          <button>
+            /about
+          </button>
+        </div>"
+      `);
+
+      fireEvent.click(screen.getByText("/about"));
+      await waitFor(() => screen.getByText("Loading..."));
+
+      expect(getHtml(container)).toMatchInlineSnapshot(`
+        "<div>
+          <p>
+            Loading...
+          </p>
+        </div>"
+      `);
+
+      await act(() => resolve());
+      await waitFor(() => screen.getByText("About"));
+
+      expect(getHtml(container)).toMatchInlineSnapshot(`
+        "<div>
+          <h1>
+            About
+          </h1>
+        </div>"
+      `);
+    }
+
+    // eslint-disable-next-line jest/expect-expect
+    it("MemoryRouter", async () => {
+      let { Home, About, resolve } = getComponents(true);
+
+      let { container } = render(
+        <MemoryRouter>
+          <Routes>
+            <Route path="/" element={<Home />} />
+            <Route
+              path="/about"
+              element={
+                <React.Suspense fallback={<p>Loading...</p>}>
+                  <About />
+                </React.Suspense>
+              }
+            />
+          </Routes>
+        </MemoryRouter>
+      );
+
+      await assertNavigation(container, resolve);
+    });
+
+    // eslint-disable-next-line jest/expect-expect
+    it("BrowserRouter", async () => {
+      let { Home, About, resolve } = getComponents(true);
+
+      let { container } = render(
+        <BrowserRouter window={getWindowImpl("/", false)}>
+          <Routes>
+            <Route path="/" element={<Home />} />
+            <Route
+              path="/about"
+              element={
+                <React.Suspense fallback={<p>Loading...</p>}>
+                  <About />
+                </React.Suspense>
+              }
+            />
+          </Routes>
+        </BrowserRouter>
+      );
+
+      await assertNavigation(container, resolve);
+    });
+
+    // eslint-disable-next-line jest/expect-expect
+    it("HashRouter", async () => {
+      let { Home, About, resolve } = getComponents(true);
+
+      let { container } = render(
+        <HashRouter window={getWindowImpl("/", true)}>
+          <Routes>
+            <Route path="/" element={<Home />} />
+            <Route
+              path="/about"
+              element={
+                <React.Suspense fallback={<p>Loading...</p>}>
+                  <About />
+                </React.Suspense>
+              }
+            />
+          </Routes>
+        </HashRouter>
+      );
+
+      await assertNavigation(container, resolve);
+    });
+
+    // eslint-disable-next-line jest/expect-expect
+    it("RouterProvider", async () => {
+      let { Home, About, resolve } = getComponents(true);
+
+      let router = createMemoryRouter(
+        createRoutesFromElements(
+          <>
+            <Route path="/" element={<Home />} />
+            <Route
+              path="/about"
+              element={
+                <React.Suspense fallback={<p>Loading...</p>}>
+                  <About />
+                </React.Suspense>
+              }
+            />
+          </>
+        )
+      );
+      let { container } = render(<RouterProvider router={router} />);
+
+      await assertNavigation(container, resolve);
+    });
+  });
+
+  describe("when the destination route suspends without a boundary", () => {
+    async function assertNavigation(
+      container: HTMLElement,
+      resolve: () => void
+    ) {
+      expect(getHtml(container)).toMatchInlineSnapshot(`
+        "<div>
+          <h1>
+            Home
+          </h1>
+          <button>
+            /about
+          </button>
+        </div>"
+      `);
+
+      fireEvent.click(screen.getByText("/about"));
+      await tick();
+
+      expect(getHtml(container)).toMatchInlineSnapshot(`
+        "<div>
+          <h1>
+            Home
+          </h1>
+          <button>
+            /about
+          </button>
+        </div>"
+      `);
+
+      await act(() => resolve());
+      await waitFor(() => screen.getByText("About"));
+
+      expect(getHtml(container)).toMatchInlineSnapshot(`
+        "<div>
+          <h1>
+            About
+          </h1>
+        </div>"
+      `);
+    }
+
+    // eslint-disable-next-line jest/expect-expect
+    it("MemoryRouter", async () => {
+      let { Home, About, resolve } = getComponents(false);
+
+      let { container } = render(
+        <MemoryRouter>
+          <Routes>
+            <Route path="/" element={<Home />} />
+            <Route path="/about" element={<About />} />
+          </Routes>
+        </MemoryRouter>
+      );
+
+      await assertNavigation(container, resolve);
+    });
+
+    // eslint-disable-next-line jest/expect-expect
+    it("BrowserRouter", async () => {
+      let { Home, About, resolve } = getComponents(false);
+
+      let { container } = render(
+        <BrowserRouter window-={getWindowImpl("/", true)}>
+          <Routes>
+            <Route path="/" element={<Home />} />
+            <Route path="/about" element={<About />} />
+          </Routes>
+        </BrowserRouter>
+      );
+
+      await assertNavigation(container, resolve);
+    });
+
+    // eslint-disable-next-line jest/expect-expect
+    it("HashRouter", async () => {
+      let { Home, About, resolve } = getComponents(false);
+
+      let { container } = render(
+        <HashRouter window-={getWindowImpl("/", true)}>
+          <Routes>
+            <Route path="/" element={<Home />} />
+            <Route path="/about" element={<About />} />
+          </Routes>
+        </HashRouter>
+      );
+
+      await assertNavigation(container, resolve);
+    });
+
+    // eslint-disable-next-line jest/expect-expect
+    it("RouterProvider", async () => {
+      let { Home, About, resolve } = getComponents(false);
+
+      let router = createMemoryRouter(
+        createRoutesFromElements(
+          <>
+            <Route path="/" element={<Home />} />
+            <Route path="/about" element={<About />} />
+          </>
+        )
+      );
+      let { container } = render(<RouterProvider router={router} />);
+
+      await assertNavigation(container, resolve);
+    });
+  });
+});
+
+function getWindowImpl(initialUrl: string, isHash = false): Window {
+  // Need to use our own custom DOM in order to get a working history
+  const dom = new JSDOM(`<!DOCTYPE html>`, { url: "http://localhost/" });
+  dom.window.history.replaceState(null, "", (isHash ? "#" : "") + initialUrl);
+  return dom.window as unknown as Window;
+}
+
+function getHtml(container: HTMLElement) {
+  return prettyDOM(container, undefined, {
+    highlight: false,
+  });
+}
+
+async function tick() {
+  await new Promise((r) => setTimeout(r, 0));
+}
+
+function createDeferred() {
+  let resolve: (val?: any) => Promise<void>;
+  let reject: (error?: Error) => Promise<void>;
+  let promise = new Promise((res, rej) => {
+    resolve = async (val: any) => {
+      res(val);
+      await tick();
+      await promise;
+    };
+    reject = async (error?: Error) => {
+      rej(error);
+      await promise.catch(() => tick());
+    };
+  });
+  return {
+    promise,
+    //@ts-ignore
+    resolve,
+    //@ts-ignore
+    reject,
+  };
+}
diff --git a/packages/react-router-dom/index.tsx b/packages/react-router-dom/index.tsx
index 7dd37cf655..c323e4ca9f 100644
--- a/packages/react-router-dom/index.tsx
+++ b/packages/react-router-dom/index.tsx
@@ -4,7 +4,9 @@
  */
 import * as React from "react";
 import type {
+  Location,
   NavigateOptions,
+  NavigationType,
   RelativeRoutingType,
   RouteObject,
   To,
@@ -312,12 +314,20 @@ export function BrowserRouter({
   }
 
   let history = historyRef.current;
-  let [state, setState] = React.useState({
+  let [state, setStateImpl] = React.useState({
     action: history.action,
     location: history.location,
   });
+  let setState = React.useCallback(
+    (newState: { action: NavigationType; location: Location }) => {
+      "startTransition" in React
+        ? React.startTransition(() => setStateImpl(newState))
+        : setStateImpl(newState);
+    },
+    [setStateImpl]
+  );
 
-  React.useLayoutEffect(() => history.listen(setState), [history]);
+  React.useLayoutEffect(() => history.listen(setState), [history, setState]);
 
   return (
     <Router
@@ -347,12 +357,20 @@ export function HashRouter({ basename, children, window }: HashRouterProps) {
   }
 
   let history = historyRef.current;
-  let [state, setState] = React.useState({
+  let [state, setStateImpl] = React.useState({
     action: history.action,
     location: history.location,
   });
+  let setState = React.useCallback(
+    (newState: { action: NavigationType; location: Location }) => {
+      "startTransition" in React
+        ? React.startTransition(() => setStateImpl(newState))
+        : setStateImpl(newState);
+    },
+    [setStateImpl]
+  );
 
-  React.useLayoutEffect(() => history.listen(setState), [history]);
+  React.useLayoutEffect(() => history.listen(setState), [history, setState]);
 
   return (
     <Router
@@ -378,12 +396,20 @@ export interface HistoryRouterProps {
  * version of the history library that React Router uses internally.
  */
 function HistoryRouter({ basename, children, history }: HistoryRouterProps) {
-  const [state, setState] = React.useState({
+  let [state, setStateImpl] = React.useState({
     action: history.action,
     location: history.location,
   });
+  let setState = React.useCallback(
+    (newState: { action: NavigationType; location: Location }) => {
+      "startTransition" in React
+        ? React.startTransition(() => setStateImpl(newState))
+        : setStateImpl(newState);
+    },
+    [setStateImpl]
+  );
 
-  React.useLayoutEffect(() => history.listen(setState), [history]);
+  React.useLayoutEffect(() => history.listen(setState), [history, setState]);
 
   return (
     <Router
diff --git a/packages/react-router/lib/components.tsx b/packages/react-router/lib/components.tsx
index 56ef4c1507..184426804b 100644
--- a/packages/react-router/lib/components.tsx
+++ b/packages/react-router/lib/components.tsx
@@ -63,7 +63,15 @@ export function RouterProvider({
 }: RouterProviderProps): React.ReactElement {
   // Need to use a layout effect here so we are subscribed early enough to
   // pick up on any render-driven redirects/navigations (useEffect/<Navigate>)
-  let [state, setState] = React.useState(router.state);
+  let [state, setStateImpl] = React.useState(router.state);
+  let setState = React.useCallback(
+    (newState: RouterState) => {
+      "startTransition" in React
+        ? React.startTransition(() => setStateImpl(newState))
+        : setStateImpl(newState);
+    },
+    [setStateImpl]
+  );
   React.useLayoutEffect(() => router.subscribe(setState), [router, setState]);
 
   let navigator = React.useMemo((): Navigator => {
@@ -164,12 +172,20 @@ export function MemoryRouter({
   }
 
   let history = historyRef.current;
-  let [state, setState] = React.useState({
+  let [state, setStateImpl] = React.useState({
     action: history.action,
     location: history.location,
   });
+  let setState = React.useCallback(
+    (newState: { action: NavigationType; location: Location }) => {
+      "startTransition" in React
+        ? React.startTransition(() => setStateImpl(newState))
+        : setStateImpl(newState);
+    },
+    [setStateImpl]
+  );
 
-  React.useLayoutEffect(() => history.listen(setState), [history]);
+  React.useLayoutEffect(() => history.listen(setState), [history, setState]);
 
   return (
     <Router

From 1ebe206ff57b7bab8897dbdffc7e402802f35fb7 Mon Sep 17 00:00:00 2001
From: Matt Brophy <matt@brophy.org>
Date: Wed, 3 May 2023 16:53:24 -0400
Subject: [PATCH 2/5] Update unit test to capture concurrent mode behavior

---
 packages/react-router/__tests__/navigate-test.tsx | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/packages/react-router/__tests__/navigate-test.tsx b/packages/react-router/__tests__/navigate-test.tsx
index 5342758b5e..7c12862284 100644
--- a/packages/react-router/__tests__/navigate-test.tsx
+++ b/packages/react-router/__tests__/navigate-test.tsx
@@ -639,9 +639,6 @@ describe("<Navigate>", () => {
           {
             index: true,
             Component() {
-              // When state managed by react and changes during render, we'll
-              // only "see" the value from the first pass through here in our
-              // effects
               let [count, setCount] = React.useState(0);
               React.useEffect(() => {
                 if (count === 0) {
@@ -684,11 +681,11 @@ describe("<Navigate>", () => {
           Page B
         </h1>
         <p>
-          0
+          1
         </p>
       </div>"
     `);
-    expect(navigateSpy).toHaveBeenCalledTimes(2);
+    expect(navigateSpy).toHaveBeenCalledTimes(3);
     expect(navigateSpy.mock.calls[0]).toMatchObject([
       { pathname: "/b" },
       { state: { count: 0 } },
@@ -697,7 +694,7 @@ describe("<Navigate>", () => {
       { pathname: "/b" },
       { state: { count: 0 } },
     ]);
-    expect(renders).toEqual([0, 0]);
+    expect(renders).toEqual([1, 1]);
   });
 
   it("handles setState in render in StrictMode using a data router (async loader)", async () => {

From d5a81ff9ecdbc46ec3f2144738af6489c231c8a5 Mon Sep 17 00:00:00 2001
From: Matt Brophy <matt@brophy.org>
Date: Tue, 16 May 2023 16:32:31 -0400
Subject: [PATCH 3/5] Remove uneeded suspense boundary from test

---
 .../concurrent-mode-navigations-test.tsx      | 33 +++++++------------
 1 file changed, 11 insertions(+), 22 deletions(-)

diff --git a/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx b/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx
index ae0b1008d6..29f384eefb 100644
--- a/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx
+++ b/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx
@@ -21,7 +21,7 @@ import {
 import { JSDOM } from "jsdom";
 
 describe("Handles concurrent mode features during navigations", () => {
-  function getComponents(withBoundary: boolean) {
+  function getComponents() {
     function Home() {
       let navigate = useNavigate();
       return (
@@ -39,27 +39,16 @@ describe("Handles concurrent mode features during navigations", () => {
       return dfd.resolve();
     }
 
-    function AboutWithoutBoundary() {
+    function About() {
       if (!resolved) {
         throw dfd.promise;
       }
       return <h1>About</h1>;
     }
 
-    function AboutWithBoundary() {
-      if (!resolved) {
-        throw dfd.promise;
-      }
-      return (
-        <React.Suspense fallback={<p>Loading...</p>}>
-          <h1>About</h1>
-        </React.Suspense>
-      );
-    }
-
     return {
       Home,
-      About: withBoundary ? AboutWithBoundary : AboutWithoutBoundary,
+      About,
       resolve,
     };
   }
@@ -105,7 +94,7 @@ describe("Handles concurrent mode features during navigations", () => {
 
     // eslint-disable-next-line jest/expect-expect
     it("MemoryRouter", async () => {
-      let { Home, About, resolve } = getComponents(true);
+      let { Home, About, resolve } = getComponents();
 
       let { container } = render(
         <MemoryRouter>
@@ -128,7 +117,7 @@ describe("Handles concurrent mode features during navigations", () => {
 
     // eslint-disable-next-line jest/expect-expect
     it("BrowserRouter", async () => {
-      let { Home, About, resolve } = getComponents(true);
+      let { Home, About, resolve } = getComponents();
 
       let { container } = render(
         <BrowserRouter window={getWindowImpl("/", false)}>
@@ -151,7 +140,7 @@ describe("Handles concurrent mode features during navigations", () => {
 
     // eslint-disable-next-line jest/expect-expect
     it("HashRouter", async () => {
-      let { Home, About, resolve } = getComponents(true);
+      let { Home, About, resolve } = getComponents();
 
       let { container } = render(
         <HashRouter window={getWindowImpl("/", true)}>
@@ -174,7 +163,7 @@ describe("Handles concurrent mode features during navigations", () => {
 
     // eslint-disable-next-line jest/expect-expect
     it("RouterProvider", async () => {
-      let { Home, About, resolve } = getComponents(true);
+      let { Home, About, resolve } = getComponents();
 
       let router = createMemoryRouter(
         createRoutesFromElements(
@@ -241,7 +230,7 @@ describe("Handles concurrent mode features during navigations", () => {
 
     // eslint-disable-next-line jest/expect-expect
     it("MemoryRouter", async () => {
-      let { Home, About, resolve } = getComponents(false);
+      let { Home, About, resolve } = getComponents();
 
       let { container } = render(
         <MemoryRouter>
@@ -257,7 +246,7 @@ describe("Handles concurrent mode features during navigations", () => {
 
     // eslint-disable-next-line jest/expect-expect
     it("BrowserRouter", async () => {
-      let { Home, About, resolve } = getComponents(false);
+      let { Home, About, resolve } = getComponents();
 
       let { container } = render(
         <BrowserRouter window-={getWindowImpl("/", true)}>
@@ -273,7 +262,7 @@ describe("Handles concurrent mode features during navigations", () => {
 
     // eslint-disable-next-line jest/expect-expect
     it("HashRouter", async () => {
-      let { Home, About, resolve } = getComponents(false);
+      let { Home, About, resolve } = getComponents();
 
       let { container } = render(
         <HashRouter window-={getWindowImpl("/", true)}>
@@ -289,7 +278,7 @@ describe("Handles concurrent mode features during navigations", () => {
 
     // eslint-disable-next-line jest/expect-expect
     it("RouterProvider", async () => {
-      let { Home, About, resolve } = getComponents(false);
+      let { Home, About, resolve } = getComponents();
 
       let router = createMemoryRouter(
         createRoutesFromElements(

From d539adc56eca9b9f32b74c9d125d519f73165c28 Mon Sep 17 00:00:00 2001
From: Matt Brophy <matt@brophy.org>
Date: Wed, 17 May 2023 12:51:02 -0400
Subject: [PATCH 4/5] Bundle bump

---
 package.json | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/package.json b/package.json
index 8b95a2d15f..5e79f3be06 100644
--- a/package.json
+++ b/package.json
@@ -108,16 +108,16 @@
       "none": "45 kB"
     },
     "packages/react-router/dist/react-router.production.min.js": {
-      "none": "13.3 kB"
+      "none": "13.4 kB"
     },
     "packages/react-router/dist/umd/react-router.production.min.js": {
-      "none": "15.6 kB"
+      "none": "15.8 kB"
     },
     "packages/react-router-dom/dist/react-router-dom.production.min.js": {
-      "none": "11.8 kB"
+      "none": "12.0 kB"
     },
     "packages/react-router-dom/dist/umd/react-router-dom.production.min.js": {
-      "none": "17.7 kB"
+      "none": "17.9 kB"
     }
   }
 }

From 5c41962569ea476c66644a627d036a96e12bda72 Mon Sep 17 00:00:00 2001
From: Matt Brophy <matt@brophy.org>
Date: Wed, 17 May 2023 16:07:16 -0400
Subject: [PATCH 5/5] Add tests with React.lazy

---
 .../__tests__/components/LazyComponent.tsx    |   5 +
 .../concurrent-mode-navigations-test.tsx      | 218 +++++++++++-------
 2 files changed, 145 insertions(+), 78 deletions(-)
 create mode 100644 packages/react-router-dom/__tests__/components/LazyComponent.tsx

diff --git a/packages/react-router-dom/__tests__/components/LazyComponent.tsx b/packages/react-router-dom/__tests__/components/LazyComponent.tsx
new file mode 100644
index 0000000000..9edc993d5c
--- /dev/null
+++ b/packages/react-router-dom/__tests__/components/LazyComponent.tsx
@@ -0,0 +1,5 @@
+import * as React from "react";
+
+export default function LazyComponent() {
+  return <h1>Lazy</h1>;
+}
diff --git a/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx b/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx
index 29f384eefb..6e7839a012 100644
--- a/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx
+++ b/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx
@@ -19,6 +19,7 @@ import {
   waitFor,
 } from "@testing-library/react";
 import { JSDOM } from "jsdom";
+import LazyComponent from "./components//LazyComponent";
 
 describe("Handles concurrent mode features during navigations", () => {
   function getComponents() {
@@ -28,6 +29,7 @@ describe("Handles concurrent mode features during navigations", () => {
         <>
           <h1>Home</h1>
           <button onClick={() => navigate("/about")}>/about</button>
+          <button onClick={() => navigate("/lazy")}>/lazy</button>
         </>
       );
     }
@@ -40,61 +42,79 @@ describe("Handles concurrent mode features during navigations", () => {
     }
 
     function About() {
+      let navigate = useNavigate();
       if (!resolved) {
         throw dfd.promise;
       }
-      return <h1>About</h1>;
+      return (
+        <>
+          <h1>About</h1>
+          <button onClick={() => navigate(-1)}>back</button>
+        </>
+      );
     }
 
+    let lazyDfd = createDeferred();
+
+    const LazyComponent = React.lazy(async () => {
+      await lazyDfd.promise;
+      return import("./components/LazyComponent");
+    });
+
     return {
       Home,
       About,
+      LazyComponent,
       resolve,
+      resolveLazy: lazyDfd.resolve,
     };
   }
 
   describe("when the destination route suspends with a boundary", () => {
     async function assertNavigation(
       container: HTMLElement,
-      resolve: () => void
+      resolve: () => void,
+      resolveLazy: () => void
     ) {
-      expect(getHtml(container)).toMatchInlineSnapshot(`
-        "<div>
-          <h1>
-            Home
-          </h1>
-          <button>
-            /about
-          </button>
-        </div>"
-      `);
-
-      fireEvent.click(screen.getByText("/about"));
-      await waitFor(() => screen.getByText("Loading..."));
+      // Start on home
+      expect(getHtml(container)).toMatch("Home");
 
-      expect(getHtml(container)).toMatchInlineSnapshot(`
-        "<div>
-          <p>
-            Loading...
-          </p>
-        </div>"
-      `);
+      // Click to /about and should see Suspense boundary
+      await act(() => {
+        fireEvent.click(screen.getByText("/about"));
+      });
+      await waitFor(() => screen.getByText("Loading..."));
+      expect(getHtml(container)).toMatch("Loading...");
 
+      // Resolve the destination UI to clear the boundary
       await act(() => resolve());
       await waitFor(() => screen.getByText("About"));
-
-      expect(getHtml(container)).toMatchInlineSnapshot(`
-        "<div>
-          <h1>
-            About
-          </h1>
-        </div>"
-      `);
+      expect(getHtml(container)).toMatch("About");
+
+      // Back to home
+      await act(() => {
+        fireEvent.click(screen.getByText("back"));
+      });
+      await waitFor(() => screen.getByText("Home"));
+      expect(getHtml(container)).toMatch("Home");
+
+      // Click to /lazy and should see Suspense boundary
+      await act(() => {
+        fireEvent.click(screen.getByText("/lazy"));
+      });
+      await waitFor(() => screen.getByText("Loading Lazy Component..."));
+      expect(getHtml(container)).toMatch("Loading Lazy Component...");
+
+      // Resolve the lazy component to clear the boundary
+      await act(() => resolveLazy());
+      await waitFor(() => screen.getByText("Lazy"));
+      expect(getHtml(container)).toMatch("Lazy");
     }
 
     // eslint-disable-next-line jest/expect-expect
     it("MemoryRouter", async () => {
-      let { Home, About, resolve } = getComponents();
+      let { Home, About, LazyComponent, resolve, resolveLazy } =
+        getComponents();
 
       let { container } = render(
         <MemoryRouter>
@@ -108,16 +128,25 @@ describe("Handles concurrent mode features during navigations", () => {
                 </React.Suspense>
               }
             />
+            <Route
+              path="/lazy"
+              element={
+                <React.Suspense fallback={<p>Loading Lazy Component...</p>}>
+                  <LazyComponent />
+                </React.Suspense>
+              }
+            />
           </Routes>
         </MemoryRouter>
       );
 
-      await assertNavigation(container, resolve);
+      await assertNavigation(container, resolve, resolveLazy);
     });
 
     // eslint-disable-next-line jest/expect-expect
     it("BrowserRouter", async () => {
-      let { Home, About, resolve } = getComponents();
+      let { Home, About, LazyComponent, resolve, resolveLazy } =
+        getComponents();
 
       let { container } = render(
         <BrowserRouter window={getWindowImpl("/", false)}>
@@ -131,16 +160,25 @@ describe("Handles concurrent mode features during navigations", () => {
                 </React.Suspense>
               }
             />
+            <Route
+              path="/lazy"
+              element={
+                <React.Suspense fallback={<p>Loading Lazy Component...</p>}>
+                  <LazyComponent />
+                </React.Suspense>
+              }
+            />
           </Routes>
         </BrowserRouter>
       );
 
-      await assertNavigation(container, resolve);
+      await assertNavigation(container, resolve, resolveLazy);
     });
 
     // eslint-disable-next-line jest/expect-expect
     it("HashRouter", async () => {
-      let { Home, About, resolve } = getComponents();
+      let { Home, About, LazyComponent, resolve, resolveLazy } =
+        getComponents();
 
       let { container } = render(
         <HashRouter window={getWindowImpl("/", true)}>
@@ -154,16 +192,25 @@ describe("Handles concurrent mode features during navigations", () => {
                 </React.Suspense>
               }
             />
+            <Route
+              path="/lazy"
+              element={
+                <React.Suspense fallback={<p>Loading Lazy Component...</p>}>
+                  <LazyComponent />
+                </React.Suspense>
+              }
+            />
           </Routes>
         </HashRouter>
       );
 
-      await assertNavigation(container, resolve);
+      await assertNavigation(container, resolve, resolveLazy);
     });
 
     // eslint-disable-next-line jest/expect-expect
     it("RouterProvider", async () => {
-      let { Home, About, resolve } = getComponents();
+      let { Home, About, LazyComponent, resolve, resolveLazy } =
+        getComponents();
 
       let router = createMemoryRouter(
         createRoutesFromElements(
@@ -177,120 +224,135 @@ describe("Handles concurrent mode features during navigations", () => {
                 </React.Suspense>
               }
             />
+            <Route
+              path="/lazy"
+              element={
+                <React.Suspense fallback={<p>Loading Lazy Component...</p>}>
+                  <LazyComponent />
+                </React.Suspense>
+              }
+            />
           </>
         )
       );
       let { container } = render(<RouterProvider router={router} />);
 
-      await assertNavigation(container, resolve);
+      await assertNavigation(container, resolve, resolveLazy);
     });
   });
 
   describe("when the destination route suspends without a boundary", () => {
     async function assertNavigation(
       container: HTMLElement,
-      resolve: () => void
+      resolve: () => void,
+      resolveLazy: () => void
     ) {
-      expect(getHtml(container)).toMatchInlineSnapshot(`
-        "<div>
-          <h1>
-            Home
-          </h1>
-          <button>
-            /about
-          </button>
-        </div>"
-      `);
-
-      fireEvent.click(screen.getByText("/about"));
-      await tick();
+      // Start on home
+      expect(getHtml(container)).toMatch("Home");
 
-      expect(getHtml(container)).toMatchInlineSnapshot(`
-        "<div>
-          <h1>
-            Home
-          </h1>
-          <button>
-            /about
-          </button>
-        </div>"
-      `);
+      // Click to /about and should see the frozen current UI
+      await act(() => {
+        fireEvent.click(screen.getByText("/about"));
+      });
+      await waitFor(() => screen.getByText("Home"));
+      expect(getHtml(container)).toMatch("Home");
 
+      // Resolve the destination UI to clear the boundary
       await act(() => resolve());
       await waitFor(() => screen.getByText("About"));
-
-      expect(getHtml(container)).toMatchInlineSnapshot(`
-        "<div>
-          <h1>
-            About
-          </h1>
-        </div>"
-      `);
+      expect(getHtml(container)).toMatch("About");
+
+      // Back to home
+      await act(() => {
+        fireEvent.click(screen.getByText("back"));
+      });
+      await waitFor(() => screen.getByText("Home"));
+      expect(getHtml(container)).toMatch("Home");
+
+      // Click to /lazy and should see the frozen current UI
+      await act(() => {
+        fireEvent.click(screen.getByText("/lazy"));
+      });
+      await waitFor(() => screen.getByText("Home"));
+      expect(getHtml(container)).toMatch("Home");
+
+      // Resolve the lazy component to clear the boundary
+      await act(() => resolveLazy());
+      await waitFor(() => screen.getByText("Lazy"));
+      expect(getHtml(container)).toMatch("Lazy");
     }
 
     // eslint-disable-next-line jest/expect-expect
     it("MemoryRouter", async () => {
-      let { Home, About, resolve } = getComponents();
+      let { Home, About, resolve, LazyComponent, resolveLazy } =
+        getComponents();
 
       let { container } = render(
         <MemoryRouter>
           <Routes>
             <Route path="/" element={<Home />} />
             <Route path="/about" element={<About />} />
+            <Route path="/lazy" element={<LazyComponent />} />
           </Routes>
         </MemoryRouter>
       );
 
-      await assertNavigation(container, resolve);
+      await assertNavigation(container, resolve, resolveLazy);
     });
 
     // eslint-disable-next-line jest/expect-expect
     it("BrowserRouter", async () => {
-      let { Home, About, resolve } = getComponents();
+      let { Home, About, resolve, LazyComponent, resolveLazy } =
+        getComponents();
 
       let { container } = render(
         <BrowserRouter window-={getWindowImpl("/", true)}>
           <Routes>
             <Route path="/" element={<Home />} />
             <Route path="/about" element={<About />} />
+            <Route path="/lazy" element={<LazyComponent />} />
           </Routes>
         </BrowserRouter>
       );
 
-      await assertNavigation(container, resolve);
+      await assertNavigation(container, resolve, resolveLazy);
     });
 
     // eslint-disable-next-line jest/expect-expect
     it("HashRouter", async () => {
-      let { Home, About, resolve } = getComponents();
+      let { Home, About, resolve, LazyComponent, resolveLazy } =
+        getComponents();
 
       let { container } = render(
         <HashRouter window-={getWindowImpl("/", true)}>
           <Routes>
             <Route path="/" element={<Home />} />
             <Route path="/about" element={<About />} />
+            <Route path="/lazy" element={<LazyComponent />} />
           </Routes>
         </HashRouter>
       );
 
-      await assertNavigation(container, resolve);
+      await assertNavigation(container, resolve, resolveLazy);
     });
 
     // eslint-disable-next-line jest/expect-expect
     it("RouterProvider", async () => {
-      let { Home, About, resolve } = getComponents();
+      let { Home, About, resolve, LazyComponent, resolveLazy } =
+        getComponents();
 
       let router = createMemoryRouter(
         createRoutesFromElements(
           <>
             <Route path="/" element={<Home />} />
             <Route path="/about" element={<About />} />
+            <Route path="/lazy" element={<LazyComponent />} />
           </>
         )
       );
       let { container } = render(<RouterProvider router={router} />);
 
-      await assertNavigation(container, resolve);
+      await assertNavigation(container, resolve, resolveLazy);
     });
   });
 });