Fix component props type and update tests

src/__tests__/index.test.tsx

@@ -1,14 +1,25 @@
-import React from "react";
+import React, {
+    forwardRef,
+    ReactNode,
+    RefObject,
+    Suspense,
+    useRef,
+} from "react";
 import { act, render, screen, waitFor } from "@testing-library/react";
-import lazy, { lazyWithPreload as namedExport } from "../index";
+import lazyWithPreload, { lazyWithPreload as namedExport } from "../index";
+
+interface TestComponentProps {
+    foo: string;
+    children: ReactNode;
+}
 
 function getTestComponentModule() {
-    const TestComponent = React.forwardRef<
-        HTMLDivElement,
-        { foo: string; children: React.ReactNode }
-    >(function TestComponent(props, ref) {
-        return <div ref={ref}>{`${props.foo} ${props.children}`}</div>;
-    });
+    const TestComponent = forwardRef<HTMLDivElement, TestComponentProps>(
+        function TestComponent(props, ref) {
+            return <div ref={ref}>{`${props.foo} ${props.children}`}</div>;
+        }
+    );
+
     let loaded = false;
     let loadCalls = 0;
 
@@ -18,116 +29,141 @@ function getTestComponentModule() {
         OriginalComponent: TestComponent,
         TestComponent: async () => {
             loaded = true;
-            loadCalls++;
+            loadCalls += 1;
             return { default: TestComponent };
         },
     };
 }
 
-describe("lazy", () => {
+describe("lazyWithPreload", () => {
     it("renders normally without invoking preload", async () => {
+        // Given
         const { TestComponent, isLoaded } = getTestComponentModule();
-        const LazyTestComponent = lazy(TestComponent);
+        const LazyTestComponent = lazyWithPreload(TestComponent);
 
+        // Then
         expect(isLoaded()).toBe(false);
 
+        // When
         render(
-            <React.Suspense fallback={null}>
+            <Suspense fallback={null}>
                 <LazyTestComponent foo="bar">baz</LazyTestComponent>
-            </React.Suspense>
+            </Suspense>
         );
 
+        // Then
         await waitFor(() => expect(screen.queryByText("bar baz")).toBeTruthy());
     });
 
     it("renders normally when invoking preload", async () => {
+        // Given
         const { TestComponent, isLoaded } = getTestComponentModule();
-        const LazyTestComponent = lazy(TestComponent);
+        const LazyTestComponent = lazyWithPreload(TestComponent);
+
+        // When
         await LazyTestComponent.preload();
 
+        // Then
         expect(isLoaded()).toBe(true);
 
+        // When
         render(
-            <React.Suspense fallback={null}>
+            <Suspense fallback={null}>
                 <LazyTestComponent foo="bar">baz</LazyTestComponent>
-            </React.Suspense>
+            </Suspense>
         );
 
+        // Then
         await waitFor(() => expect(screen.queryByText("bar baz")).toBeTruthy());
     });
 
     it("never renders fallback if preloaded before first render", async () => {
+        // Given
         let fallbackRendered = false;
         const Fallback = () => {
             fallbackRendered = true;
             return null;
         };
         const { TestComponent } = getTestComponentModule();
-        const LazyTestComponent = lazy(TestComponent);
-        await LazyTestComponent.preload();
+        const LazyTestComponent = lazyWithPreload(TestComponent);
 
+        // When
+        await LazyTestComponent.preload();
         render(
-            <React.Suspense fallback={<Fallback />}>
+            <Suspense fallback={<Fallback />}>
                 <LazyTestComponent foo="bar">baz</LazyTestComponent>
-            </React.Suspense>
+            </Suspense>
         );
 
+        // Then
         expect(fallbackRendered).toBe(false);
 
+        // Post
         await LazyTestComponent.preload();
     });
 
     it("renders fallback if not preloaded", async () => {
+        // Given
         let fallbackRendered = false;
         const Fallback = () => {
             fallbackRendered = true;
             return null;
         };
         const { TestComponent } = getTestComponentModule();
-        const LazyTestComponent = lazy(TestComponent);
+        const LazyTestComponent = lazyWithPreload(TestComponent);
 
+        // When
         render(
-            <React.Suspense fallback={<Fallback />}>
+            <Suspense fallback={<Fallback />}>
                 <LazyTestComponent foo="bar">baz</LazyTestComponent>
-            </React.Suspense>
+            </Suspense>
         );
 
+        // Then
         expect(fallbackRendered).toBe(true);
 
+        // Post
         await act(async () => {
             await LazyTestComponent.preload();
         });
     });
 
     it("only preloads once when preload is invoked multiple times", async () => {
+        // Given
         const { TestComponent, loadCalls } = getTestComponentModule();
-        const LazyTestComponent = lazy(TestComponent);
+        const LazyTestComponent = lazyWithPreload(TestComponent);
+
+        // When
         const preloadPromise1 = LazyTestComponent.preload();
         const preloadPromise2 = LazyTestComponent.preload();
 
         await Promise.all([preloadPromise1, preloadPromise2]);
 
+        // Then
         // If `preload()` called multiple times, it should return the same promise
         expect(preloadPromise1).toBe(preloadPromise2);
         expect(loadCalls()).toBe(1);
 
+        // When
         render(
-            <React.Suspense fallback={null}>
+            <Suspense fallback={null}>
                 <LazyTestComponent foo="bar">baz</LazyTestComponent>
-            </React.Suspense>
+            </Suspense>
         );
 
+        // Then
         await waitFor(() => expect(screen.queryByText("bar baz")).toBeTruthy());
     });
 
     it("supports ref forwarding", async () => {
+        // Given
         const { TestComponent } = getTestComponentModule();
-        const LazyTestComponent = lazy(TestComponent);
+        const LazyTestComponent = lazyWithPreload(TestComponent);
 
-        let ref: React.RefObject<HTMLDivElement> | undefined;
+        let ref: RefObject<HTMLDivElement> | undefined;
 
         function ParentComponent() {
-            ref = React.useRef<HTMLDivElement>(null);
+            ref = useRef<HTMLDivElement>(null);
 
             return (
                 <LazyTestComponent foo="bar" ref={ref}>
@@ -136,26 +172,32 @@ describe("lazy", () => {
             );
         }
 
+        // When
         render(
-            <React.Suspense fallback={null}>
+            <Suspense fallback={null}>
                 <ParentComponent />
-            </React.Suspense>
+            </Suspense>
         );
 
+        // Then
         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 () => {
+        // Given
         const { TestComponent, OriginalComponent } = getTestComponentModule();
-        const LazyTestComponent = lazy(TestComponent);
+        const LazyTestComponent = lazyWithPreload(TestComponent);
 
+        // When
         const preloadedComponent = await LazyTestComponent.preload();
 
+        // Then
         expect(preloadedComponent).toBe(OriginalComponent);
     });
 
     it("exports named export as well", () => {
-        expect(lazy).toBe(namedExport);
+        // Then
+        expect(lazyWithPreload).toBe(namedExport);
     });
 });

src/index.ts

@@ -1,24 +1,63 @@
-import { ComponentType, createElement, forwardRef, lazy } from "react";
+import {
+    ComponentProps,
+    ComponentType,
+    createElement,
+    forwardRef,
+    FunctionComponent,
+    lazy,
+    LazyExoticComponent,
+} from "react";
 
-export type PreloadableComponent<T extends ComponentType<any>> = T & {
-    preload: () => Promise<T>;
+export type PreloadableComponent<
+    COMPONENT extends ComponentType<ComponentProps<COMPONENT>>
+> = LazyExoticComponent<COMPONENT> & {
+    preload(): Promise<COMPONENT>;
 };
 
-export function lazyWithPreload<T extends ComponentType<any>>(
-    factory: () => Promise<{ default: T }>
-): PreloadableComponent<T> {
+/**
+ * Function wraps the `React.lazy()` API and adds the ability to preload the component
+ * before it is rendered for the first time.
+ *
+ * @example
+ * ```tsx
+ * import React, { Suspense } from 'react';
+ * import { lazyWithPreload } from 'react-lazy-with-preload';
+ *
+ * const LazyComponent = lazyWithPreload(() => import('./LazyComponent'));
+ *
+ * function SomeComponent() {
+ *   return (
+ *     <Suspense fallback="Loading...">
+ *       <Link
+ *         href="..."
+ *         // This component will be needed soon. Let's preload it!
+ *         onMouseOver={() => LazyComponent.preload()}
+ *       >
+ *         Click me to navigate to page with LazyComponent
+ *       </Link>
+ *     </Suspense>
+ *   );
+ * }
+ * ```
+ */
+export function lazyWithPreload<
+    COMPONENT extends ComponentType<ComponentProps<COMPONENT>>
+>(
+    factory: () => Promise<{ default: COMPONENT }>
+): PreloadableComponent<COMPONENT> {
     const LazyComponent = lazy(factory);
-    let factoryPromise: Promise<T> | undefined;
-    let LoadedComponent: T | undefined;
+    let factoryPromise: Promise<COMPONENT> | undefined;
+    let LoadedComponent: COMPONENT | undefined;
 
     const Component = forwardRef(function LazyWithPreload(props, ref) {
         return createElement(
-            LoadedComponent ?? LazyComponent,
-            Object.assign(ref ? { ref } : {}, props) as any
+            (LoadedComponent ?? LazyComponent) as FunctionComponent,
+            // eslint-disable-next-line @typescript-eslint/ban-types
+            Object.assign<{}, {}>(ref ? { ref } : {}, props)
         );
-    }) as any as PreloadableComponent<T>;
+    }) as PreloadableComponent<COMPONENT>;
 
-    Component.preload = () => {
+    Component.preload = function preload() {
         if (!factoryPromise) {
             factoryPromise = factory().then((module) => {
                 LoadedComponent = module.default;