Merge pull request #21 from yurisldk/feature/withSuspense

Add withSuspense HOC with Enhanced Type Safety, Fallback Options, and Unit Tests

src/shared/lib/react/react.hoc.tsx

@@ -3,32 +3,55 @@ import {
   Suspense,
   forwardRef,
   ComponentType,
-  ForwardedRef,
   createElement,
+  ReactNode,
+  ComponentRef,
+  ComponentProps,
+  ForwardRefExoticComponent,
+  PropsWithoutRef,
+  RefAttributes,
 } from 'react'
 
-export function withSuspense<Props extends object>(
-  component: ComponentType<Props>,
-  suspenseProps: SuspenseProps & {
-    FallbackComponent?: ComponentType
-  },
-) {
-  const Wrapped = forwardRef<ComponentType<Props>, Props>(
-    (props: Props, ref: ForwardedRef<ComponentType<Props>>) =>
-      createElement(
+type SuspenseSharedProps = Omit<SuspenseProps, 'fallback'>
+
+type SuspensePropsWithComponent = SuspenseSharedProps & {
+  fallback?: never
+  FallbackComponent: ComponentType
+}
+
+type SuspensePropsWithFallback = SuspenseSharedProps & {
+  fallback: ReactNode
+  FallbackComponent?: never
+}
+
+type WithSuspenseProps = SuspensePropsWithComponent | SuspensePropsWithFallback
+
+export function withSuspense<T extends ComponentType<any>>(
+  WrappedComponent: T,
+  suspenseProps: WithSuspenseProps,
+): ForwardRefExoticComponent<
+  PropsWithoutRef<ComponentProps<T>> & RefAttributes<ComponentRef<T>>
+> {
+  const Wrapped = forwardRef<ComponentRef<T>, ComponentProps<T>>(
+    (props, ref) => {
+      const { fallback, FallbackComponent, ...otherSuspenseProps } =
+        suspenseProps
+
+      const suspenseFallback =
+        fallback ??
+        (FallbackComponent ? createElement(FallbackComponent) : null)
+
+      return createElement(
         Suspense,
-        {
-          fallback:
-            suspenseProps.fallback ||
-            (suspenseProps.FallbackComponent &&
-              createElement(suspenseProps.FallbackComponent)),
-        },
-        createElement(component, { ...props, ref }),
-      ),
+        { ...otherSuspenseProps, fallback: suspenseFallback },
+        createElement(WrappedComponent, { ...props, ref }),
+      )
+    },
   )
 
-  const name = component.displayName || component.name || 'Unknown'
-  Wrapped.displayName = `withSuspense(${name})`
+  const wrappedComponentName =
+    WrappedComponent.displayName || WrappedComponent.name || 'Unknown'
+  Wrapped.displayName = `withSuspense(${wrappedComponentName})`
 
   return Wrapped
 }

src/shared/lib/react/react.test.tsx

@@ -0,0 +1,86 @@
+import React, {
+  ComponentPropsWithoutRef,
+  ForwardedRef,
+  forwardRef,
+  lazy,
+} from 'react'
+import { render, screen, waitFor } from '@testing-library/react'
+import { describe, expect, it } from 'vitest'
+import { withSuspense } from './react.hoc'
+
+describe('withSuspense HOC', () => {
+  it('renders the synchronous component correctly without fallback', async () => {
+    render(<SyncComponentWithSuspense />)
+
+    await waitFor(() =>
+      expect(screen.getByTestId(MOCK_COMPONENT_TEST_ID)).toBeInTheDocument(),
+    )
+  })
+
+  it('renders the fallback when the lazy component is suspended (JSX fallback)', async () => {
+    render(<LazyComponentWithJSXFallback />)
+
+    await waitFor(() =>
+      expect(
+        screen.getByTestId(FALLBACK_COMPONENT_TEST_ID),
+      ).toBeInTheDocument(),
+    )
+  })
+
+  it('renders the FallbackComponent when the lazy component is suspended (Component fallback)', async () => {
+    render(<LazyComponentWithFallbackComponent />)
+
+    await waitFor(() =>
+      expect(
+        screen.getByTestId(FALLBACK_COMPONENT_TEST_ID),
+      ).toBeInTheDocument(),
+    )
+  })
+
+  it('forwards the ref to the wrapped component', async () => {
+    const ref = React.createRef<HTMLDivElement>()
+
+    render(<SyncComponentWithSuspense ref={ref} />)
+
+    await waitFor(() =>
+      expect(screen.getByTestId(MOCK_COMPONENT_TEST_ID)).toBeInTheDocument(),
+    )
+
+    expect(ref.current).toBeInTheDocument()
+  })
+})
+
+const MOCK_COMPONENT_TEST_ID = 'mock-component'
+const FALLBACK_COMPONENT_TEST_ID = 'fallback-component'
+
+const FallbackComponent = () => <div data-testid={FALLBACK_COMPONENT_TEST_ID} />
+
+const MockComponent = forwardRef(
+  (
+    props: ComponentPropsWithoutRef<'div'>,
+    ref?: ForwardedRef<HTMLDivElement>,
+  ) => (
+    <div
+      ref={ref}
+      {...props}
+      data-testid={MOCK_COMPONENT_TEST_ID}
+    />
+  ),
+)
+
+const LazyMockComponent = lazy(
+  () =>
+    new Promise<{ default: React.ComponentType }>((resolve) => {
+      setTimeout(() => resolve({ default: MockComponent }), 500)
+    }),
+)
+
+const SyncComponentWithSuspense = withSuspense(MockComponent, {
+  FallbackComponent,
+})
+const LazyComponentWithFallbackComponent = withSuspense(LazyMockComponent, {
+  FallbackComponent,
+})
+const LazyComponentWithJSXFallback = withSuspense(LazyMockComponent, {
+  fallback: <FallbackComponent />,
+})