fix: move fallbackElement into router context

packages/react-router-dom/__tests__/data-browser-router-test.tsx

@@ -31,6 +31,7 @@ import {
   useFetchers,
   UNSAFE_DataRouterStateContext as DataRouterStateContext,
   defer,
+  useLocation,
 } from "react-router-dom";
 
 // Private API
@@ -312,6 +313,54 @@ function testDomRouter(
       `);
     });
 
+    it("renders fallbackElement within router contexts", async () => {
+      let fooDefer = createDeferred();
+      let { container } = render(
+        <TestDataRouter
+          window={getWindow("/foo")}
+          fallbackElement={<FallbackElement />}
+        >
+          <Route path="/" element={<Outlet />}>
+            <Route
+              path="foo"
+              loader={() => fooDefer.promise}
+              element={<Foo />}
+            />
+          </Route>
+        </TestDataRouter>
+      );
+
+      function FallbackElement() {
+        let location = useLocation();
+        return <p>Loading{location.pathname}</p>;
+      }
+
+      function Foo() {
+        let data = useLoaderData();
+        return <h1>Foo:{data?.message}</h1>;
+      }
+
+      expect(getHtml(container)).toMatchInlineSnapshot(`
+        "<div>
+          <p>
+            Loading
+            /foo
+          </p>
+        </div>"
+      `);
+
+      fooDefer.resolve({ message: "From Foo Loader" });
+      await waitFor(() => screen.getByText("Foo:From Foo Loader"));
+      expect(getHtml(container)).toMatchInlineSnapshot(`
+        "<div>
+          <h1>
+            Foo:
+            From Foo Loader
+          </h1>
+        </div>"
+      `);
+    });
+
     it("handles link navigations", async () => {
       render(
         <TestDataRouter window={getWindow("/foo")} hydrationData={{}}>

packages/react-router-dom/index.tsx

@@ -234,12 +234,8 @@ export function DataBrowserRouter({
   let router = routerSingleton;
 
   return (
-    <DataRouterProvider
-      router={router}
-      basename={basename}
-      fallbackElement={fallbackElement}
-    >
-      <DataRouter />
+    <DataRouterProvider router={router} basename={basename}>
+      <DataRouter fallbackElement={fallbackElement} />
     </DataRouterProvider>
   );
 }
@@ -274,12 +270,8 @@ export function DataHashRouter({
   let router = routerSingleton;
 
   return (
-    <DataRouterProvider
-      router={router}
-      basename={basename}
-      fallbackElement={fallbackElement}
-    >
-      <DataRouter />
+    <DataRouterProvider router={router} basename={basename}>
+      <DataRouter fallbackElement={fallbackElement} />
     </DataRouterProvider>
   );
 }

packages/react-router/__tests__/data-memory-router-test.tsx

@@ -11,11 +11,7 @@ import "@testing-library/jest-dom";
 import type { FormMethod, Router } from "@remix-run/router";
 import { createMemoryRouter } from "@remix-run/router";
 
-import {
-  DataMemoryRouterProps,
-  UNSAFE_DataRouter as DataRouter,
-  UNSAFE_DataRouterProvider as DataRouterProvider,
-} from "react-router";
+import type { DataMemoryRouterProps } from "react-router";
 import {
   DataMemoryRouter,
   Await,
@@ -26,14 +22,17 @@ import {
   useAsyncError,
   useAsyncValue,
   useLoaderData,
+  useLocation,
   useMatches,
   useRouteLoaderData,
   useRouteError,
   useNavigation,
   useRevalidator,
   MemoryRouter,
   Routes,
+  UNSAFE_DataRouter as DataRouter,
   UNSAFE_DataRouterContext as DataRouterContext,
+  UNSAFE_DataRouterProvider as DataRouterProvider,
 } from "react-router";
 
 // Private API
@@ -301,6 +300,54 @@ describe("<DataMemoryRouter>", () => {
     `);
   });
 
+  it("renders fallbackElement within router contexts", async () => {
+    let fooDefer = createDeferred();
+    let { container } = render(
+      <DataMemoryRouter
+        fallbackElement={<FallbackElement />}
+        initialEntries={["/foo"]}
+      >
+        <Route path="/" element={<Outlet />}>
+          <Route path="foo" loader={() => fooDefer.promise} element={<Foo />} />
+        </Route>
+      </DataMemoryRouter>
+    );
+
+    function FallbackElement() {
+      let location = useLocation();
+      return (
+        <>
+          <p>Loading{location.pathname}</p>
+        </>
+      );
+    }
+
+    function Foo() {
+      let data = useLoaderData();
+      return <h1>Foo:{data?.message}</h1>;
+    }
+
+    expect(getHtml(container)).toMatchInlineSnapshot(`
+      "<div>
+        <p>
+          Loading
+          /foo
+        </p>
+      </div>"
+    `);
+
+    fooDefer.resolve({ message: "From Foo Loader" });
+    await waitFor(() => screen.getByText("Foo:From Foo Loader"));
+    expect(getHtml(container)).toMatchInlineSnapshot(`
+      "<div>
+        <h1>
+          Foo:
+          From Foo Loader
+        </h1>
+      </div>"
+    `);
+  });
+
   it("handles link navigations", async () => {
     render(
       <DataMemoryRouter initialEntries={["/foo"]} hydrationData={{}}>

packages/react-router/lib/components.tsx

@@ -58,21 +58,21 @@ export function _resetModuleScope() {
   routerSingleton = null;
 }
 
+interface DataRouterProviderProps {
+  basename?: string;
+  children?: React.ReactNode;
+  router: RemixRouter;
+}
+
 /**
  * A higher-order component that, given a Remix Router instance. setups the
  * Context's required for data routing
  */
 export function DataRouterProvider({
   basename,
   children,
-  fallbackElement,
   router,
-}: {
-  basename?: string;
-  children?: React.ReactNode;
-  fallbackElement?: React.ReactNode;
-  router: RemixRouter;
-}): React.ReactElement {
+}: DataRouterProviderProps): React.ReactElement {
   // Sync router state to our component state to force re-renders
   let state: RouterState = useSyncExternalStoreShim(
     router.subscribe,
@@ -98,29 +98,29 @@ export function DataRouterProvider({
     };
   }, [router]);
 
-  let dataRouterContext: DataRouterContextObject = {
-    router,
-    navigator,
-    static: false,
-    basename: basename || "/",
-  };
-
-  if (!state.initialized) {
-    return <>{fallbackElement}</>;
-  }
-
   return (
-    <DataRouterContext.Provider value={dataRouterContext}>
+    <DataRouterContext.Provider
+      value={{
+        router,
+        navigator,
+        static: false,
+        basename: basename || "/",
+      }}
+    >
       <DataRouterStateContext.Provider value={state} children={children} />
     </DataRouterContext.Provider>
   );
 }
 
+interface DataRouterProps {
+  fallbackElement?: React.ReactNode;
+}
+
 /**
  * A data-aware wrapper for `<Router>` that leverages the Context's provided by
  * `<DataRouterProvider>`
  */
-export function DataRouter() {
+export function DataRouter({ fallbackElement }: DataRouterProps) {
   let dataRouterContext = React.useContext(DataRouterContext);
   invariant(
     dataRouterContext,
@@ -135,7 +135,7 @@ export function DataRouter() {
       navigationType={router.state.historyAction}
       navigator={navigator}
     >
-      <Routes />
+      {router.state.initialized ? <Routes /> : fallbackElement}
     </Router>
   );
 }
@@ -173,12 +173,8 @@ export function DataMemoryRouter({
   let router = routerSingleton;
 
   return (
-    <DataRouterProvider
-      router={router}
-      basename={basename}
-      fallbackElement={fallbackElement}
-    >
-      <DataRouter />
+    <DataRouterProvider router={router} basename={basename}>
+      <DataRouter fallbackElement={fallbackElement} />
     </DataRouterProvider>
   );
 }