instrument lazy - functioni version w/type errors still

Author: brophdawg11

packages/react-router/__tests__/router/instrumentation-test.ts

@@ -1,12 +1,52 @@
+import { LoaderFunction } from "../../lib/router/utils";
 import { cleanup, setup } from "./utils/data-router-setup";
-import { createFormData } from "./utils/utils";
+import { createDeferred, createFormData, tick } from "./utils/utils";
 
 // Detect any failures inside the router navigate code
 afterEach(() => {
   cleanup();
 });
 
 describe("instrumentation", () => {
+  it("allows instrumentation of lazy", async () => {
+    let spy = jest.fn();
+    let lazyDfd = createDeferred<{ loader: LoaderFunction }>();
+    let t = setup({
+      routes: [
+        {
+          index: true,
+        },
+        {
+          id: "page",
+          path: "/page",
+          lazy: () => lazyDfd.promise,
+        },
+      ],
+      unstable_instrumentRoute: (route) => {
+        route.instrument({
+          async lazy(lazy) {
+            spy("start");
+            await lazy();
+            spy("end");
+          },
+        });
+      },
+    });
+
+    await t.navigate("/page");
+    expect(spy.mock.calls).toEqual([["start"]]);
+
+    await lazyDfd.resolve({ loader: () => "PAGE" });
+    expect(spy.mock.calls).toEqual([["start"], ["end"]]);
+    await tick();
+    expect(spy.mock.calls).toEqual([["start"], ["end"]]);
+    expect(t.router.state).toMatchObject({
+      navigation: { state: "idle" },
+      location: { pathname: "/page" },
+      loaderData: { page: "PAGE" },
+    });
+  });
+
   it("allows instrumentation of loaders", async () => {
     let spy = jest.fn();
     let t = setup({
@@ -80,6 +120,107 @@ describe("instrumentation", () => {
     });
   });
 
+  it("allows instrumentation of loaders when lazy is used", async () => {
+    let spy = jest.fn();
+    let lazyDfd = createDeferred<{ loader: LoaderFunction }>();
+    let loaderDfd = createDeferred();
+    let t = setup({
+      routes: [
+        {
+          index: true,
+        },
+        {
+          id: "page",
+          path: "/page",
+          lazy: () => lazyDfd.promise,
+        },
+      ],
+      unstable_instrumentRoute: (route) => {
+        route.instrument({
+          async loader(loader) {
+            spy("start");
+            await loader();
+            spy("end");
+          },
+        });
+      },
+    });
+
+    await t.navigate("/page");
+    expect(spy).not.toHaveBeenCalled();
+
+    await lazyDfd.resolve({ loader: () => loaderDfd.promise });
+    expect(spy.mock.calls).toEqual([["start"]]);
+
+    await loaderDfd.resolve("PAGE");
+    expect(spy.mock.calls).toEqual([["start"], ["end"]]);
+
+    await tick();
+    expect(t.router.state).toMatchObject({
+      navigation: { state: "idle" },
+      location: { pathname: "/page" },
+      loaderData: { page: "PAGE" },
+    });
+  });
+
+  it("does not double-instrument when a static `loader` is used alongside `lazy`", async () => {
+    let spy = jest.fn();
+    let lazyDfd = createDeferred<{ loader: LoaderFunction }>();
+    let warnSpy = jest.spyOn(console, "warn").mockImplementation(() => {});
+    let t = setup({
+      routes: [
+        {
+          index: true,
+        },
+        {
+          id: "page",
+          path: "/page",
+          loader: true,
+          lazy: () => lazyDfd.promise,
+        },
+      ],
+      unstable_instrumentRoute: (route) => {
+        route.instrument({
+          async loader(loader) {
+            spy("start");
+            await loader();
+            spy("end");
+          },
+        });
+      },
+    });
+
+    let A = await t.navigate("/page");
+    expect(spy.mock.calls).toEqual([["start"]]);
+    await lazyDfd.resolve({ action: () => "ACTION", loader: () => "WRONG" });
+    await A.loaders.page.resolve("PAGE");
+    await tick();
+    expect(spy.mock.calls).toEqual([["start"], ["end"]]);
+    expect(t.router.state).toMatchObject({
+      navigation: { state: "idle" },
+      location: { pathname: "/page" },
+      loaderData: { page: "PAGE" },
+    });
+    spy.mockClear();
+
+    let B = await t.navigate("/page", {
+      formMethod: "POST",
+      formData: createFormData({}),
+    });
+    expect(spy).not.toHaveBeenCalled();
+    await B.loaders.page.resolve("PAGE");
+    await tick();
+    expect(spy.mock.calls).toEqual([["start"], ["end"]]);
+    expect(t.router.state).toMatchObject({
+      navigation: { state: "idle" },
+      location: { pathname: "/page" },
+      loaderData: { page: "PAGE" },
+    });
+    spy.mockClear();
+
+    warnSpy.mockRestore();
+  });
+
   it("provides read-only information to instrumentation wrappers", async () => {
     let spy = jest.fn();
     let t = setup({

packages/react-router/lib/router/instrumentation.ts

@@ -1,8 +1,10 @@
 import type {
   ActionFunction,
   AgnosticDataRouteObject,
+  LazyRouteFunction,
   LoaderFunction,
   LoaderFunctionArgs,
+  MaybePromise,
   RouterContextProvider,
 } from "./utils";
 
@@ -23,7 +25,13 @@ type InstrumentHandlerFunction = (
   info: InstrumentationInfo,
 ) => MaybePromise<void>;
 
+type InstrumentLazyFunction = (
+  handler: () => undefined,
+  info: InstrumentationInfo,
+) => MaybePromise<void>;
+
 type Instrumentations = {
+  lazy?: InstrumentLazyFunction;
   loader?: InstrumentHandlerFunction;
   action?: InstrumentHandlerFunction;
 };
@@ -35,14 +43,18 @@ type InstrumentableRoute = {
   instrument(instrumentations: Instrumentations): void;
 };
 
+const UninstrumentedSymbol = Symbol("Uninstrumented");
+
 export type unstable_InstrumentRouteFunction = (
   route: InstrumentableRoute,
 ) => void;
 
-function getInstrumentedHandler<H extends LoaderFunction | ActionFunction>(
-  impls: InstrumentHandlerFunction[],
-  handler: H,
-) {
+function getInstrumentedHandler<
+  H extends
+    | LazyRouteFunction<AgnosticDataRouteObject>
+    | LoaderFunction
+    | ActionFunction,
+>(impls: InstrumentHandlerFunction[], handler: H): H | null {
   if (impls.length === 0) {
     return null;
   }
@@ -54,7 +66,8 @@ function getInstrumentedHandler<H extends LoaderFunction | ActionFunction>(
         value = await handler(...args);
       },
     ) as unknown as (info: InstrumentationInfo) => Promise<void>;
-    await composed(getInstrumentationInfo(args[0]));
+    let composedArgs = args[0] ? [getInstrumentationInfo(args[0])] : [];
+    await composed(...composedArgs);
     return value;
   };
 }
@@ -84,10 +97,6 @@ export function getInstrumentationUpdates(
   unstable_instrumentRoute: unstable_InstrumentRouteFunction,
   route: AgnosticDataRouteObject,
 ) {
-  let updates: {
-    loader?: LoaderFunction;
-    action?: ActionFunction;
-  } = {};
   let instrumentations: Instrumentations[] = [];
   unstable_instrumentRoute({
     id: route.id,
@@ -97,18 +106,43 @@ export function getInstrumentationUpdates(
       instrumentations.push(i);
     },
   });
+
+  let updates: {
+    loader?: AgnosticDataRouteObject["loader"];
+    action?: AgnosticDataRouteObject["action"];
+    lazy?: AgnosticDataRouteObject["lazy"];
+  } = {};
+
   if (instrumentations.length > 0) {
+    if (typeof route.lazy === "function") {
+      let instrumented = getInstrumentedHandler(
+        instrumentations
+          .map((i) => i.lazy)
+          .filter(Boolean) as InstrumentHandlerFunction[],
+        route.lazy,
+      );
+      if (instrumented) {
+        updates.lazy = instrumented;
+      }
+    }
+
     if (typeof route.loader === "function") {
+      // @ts-expect-error
+      let original = route.loader[UninstrumentedSymbol] ?? route.loader;
       let instrumented = getInstrumentedHandler(
         instrumentations
           .map((i) => i.loader)
           .filter(Boolean) as InstrumentHandlerFunction[],
-        route.loader,
+        original,
       );
       if (instrumented) {
+        // @ts-expect-error Avoid double-instrumentation on lazy calls to
+        // `mapRouteProperties`
+        instrumented[UninstrumentedSymbol] = original;
         updates.loader = instrumented;
       }
     }
+
     if (typeof route.action === "function") {
       let instrumented = getInstrumentedHandler(
         instrumentations

packages/react-router/lib/router/router.ts

@@ -9,10 +9,8 @@ import {
   parsePath,
   warning,
 } from "./history";
-import {
-  getInstrumentationUpdates,
-  unstable_InstrumentRouteFunction,
-} from "./instrumentation";
+import type { unstable_InstrumentRouteFunction } from "./instrumentation";
+import { getInstrumentationUpdates } from "./instrumentation";
 import type {
   AgnosticDataRouteMatch,
   AgnosticDataRouteObject,