Return component from preload() (#15)

Author: solomonhawk

src/__tests__/index.test.tsx

@@ -15,6 +15,7 @@ function getTestComponentModule() {
     return {
         isLoaded: () => loaded,
         loadCalls: () => loadCalls,
+        OriginalComponent: TestComponent,
         TestComponent: async () => {
             loaded = true;
             loadCalls++;
@@ -138,4 +139,13 @@ describe("lazy", () => {
         await waitFor(() => expect(screen.queryByText("bar baz")).toBeTruthy());
         expect(ref?.current?.textContent).toBe("bar baz");
     });
+
+    it("returns the preloaded component when the preload promise resolves", async () => {
+        const { TestComponent, OriginalComponent } = getTestComponentModule();
+        const LazyTestComponent = lazy(TestComponent);
+
+        const preloadedComponent = await LazyTestComponent.preload()
+
+        expect(preloadedComponent).toBe(OriginalComponent);
+    });
 });

src/index.ts

@@ -1,14 +1,14 @@
 import { ComponentType, createElement, forwardRef, lazy } from "react";
 
 export type PreloadableComponent<T extends ComponentType<any>> = T & {
-    preload: () => Promise<void>;
+    preload: () => Promise<T>;
 };
 
 export default function lazyWithPreload<T extends ComponentType<any>>(
     factory: () => Promise<{ default: T }>
 ): PreloadableComponent<T> {
     const LazyComponent = lazy(factory);
-    let factoryPromise: Promise<void> | undefined;
+    let factoryPromise: Promise<T> | undefined;
     let LoadedComponent: T | undefined;
 
     const Component = (forwardRef(function LazyWithPreload(props, ref) {
@@ -22,6 +22,7 @@ export default function lazyWithPreload<T extends ComponentType<any>>(
         if (!factoryPromise) {
             factoryPromise = factory().then((module) => {
                 LoadedComponent = module.default;
+                return LoadedComponent;
             });
         }