Add support for Suspense and lazy

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@webkom/react-prepare",
-  "version": "1.1.0",
+  "version": "1.1.1",
   "description": "Prepare you app state for async server-side rendering and more!",
   "type": "module",
   "main": "./dist/react-prepare.umd.cjs",

src/prepare.ts

@@ -25,6 +25,7 @@ import {
 import {
   ConsumerElement,
   ForwardRefElement,
+  LazyElement,
   MemoElement,
   ProviderElement,
 } from './utils/reactInternalTypes';
@@ -165,8 +166,28 @@ async function prepareElement(
       const instance = createCompositeElementInstance(classElement, context);
       return [...renderCompositeElementInstance(instance, context)];
     }
+    case ELEMENT_TYPE.LAZY: {
+      const lazyElement = element as LazyElement;
+      const payload = lazyElement.type._payload;
+      const init = lazyElement.type._init;
+      while (payload._status <= 0) {
+        try {
+          await init(payload);
+        } catch (error) {
+          await error;
+        }
+      }
+      const { default: loadedElement } = payload._result;
+
+      return prepareElement(
+        { ...lazyElement, type: loadedElement },
+        errorHandler,
+        context,
+        dispatcher,
+      );
+    }
     default: {
-      throw new Error(`Unsupported element type: ${element}`);
+      throw new Error(`Unsupported element type: ${getElementType(element)}`);
     }
   }
 }

src/tests/prepare.spec.jsx

@@ -5,7 +5,9 @@ import assert from 'assert/strict';
 import sinon from 'sinon';
 import React, {
   forwardRef,
+  lazy,
   memo,
+  Suspense,
   useCallback,
   useContext,
   useDeferredValue,
@@ -17,8 +19,8 @@ import React, {
   useRef,
   useState,
   useSyncExternalStore,
-  useTransition
-} from "react";
+  useTransition,
+} from 'react';
 import PropTypes from 'prop-types';
 import { renderToStaticMarkup } from 'react-dom/server';
 import prepare from '../prepare';
@@ -27,10 +29,13 @@ import { usePreparedEffect, withPreparedEffect } from '../index';
 describe('prepare', () => {
   let originalDispatcher;
   beforeAll(() => {
-    originalDispatcher = React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.ReactCurrentDispatcher.current;
+    originalDispatcher =
+      React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+        .ReactCurrentDispatcher.current;
   });
   beforeEach(() => {
-    React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.ReactCurrentDispatcher.current = originalDispatcher;
+    React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED.ReactCurrentDispatcher.current =
+      originalDispatcher;
   });
 
   it('sets instance properties', async () => {
@@ -616,11 +621,13 @@ describe('prepare', () => {
     // eslint-disable-next-line react/display-name
     const ForwardRefComponent = forwardRef((props, ref) => {
       readContext = useContext(MyContext);
-      return (
-        <div ref={ref} />
-      );
+      return <div ref={ref} />;
     });
-    await prepare(<MyContext.Provider value="context"><ForwardRefComponent text="foo" /></MyContext.Provider>);
+    await prepare(
+      <MyContext.Provider value="context">
+        <ForwardRefComponent text="foo" />
+      </MyContext.Provider>,
+    );
     expect(readContext).toBe('context');
   });
 
@@ -690,6 +697,41 @@ describe('prepare', () => {
     );
   });
 
+  it('Should support lazy and Suspense', async () => {
+    const LazyComponent = lazy(() =>
+      Promise.resolve({ default: () => <div>Lazy</div> }),
+    );
+    const SuspenseComponent = () => (
+      <Suspense fallback={<div>Loading...</div>}>
+        <LazyComponent />
+      </Suspense>
+    );
+    await prepare(<SuspenseComponent />);
+  });
+
+  it('Should execute prepared effects inside lazy loaded components with Suspense', async () => {
+    const doAsyncSideEffect = sinon.spy(async () => {});
+    const EffectComponent = withPreparedEffect(
+      'effect',
+      doAsyncSideEffect,
+    )(() => <div>Effect</div>);
+
+    const LazyComponent = lazy(() =>
+      Promise.resolve({ default: EffectComponent }),
+    );
+    const SuspenseComponent = () => (
+      <Suspense fallback={<div>Loading...</div>}>
+        <LazyComponent prop="test" />
+      </Suspense>
+    );
+    await prepare(<SuspenseComponent />);
+    assert(doAsyncSideEffect.calledOnce, 'Should execute prepared effect');
+    assert(
+      doAsyncSideEffect.calledWith({ prop: 'test' }),
+      'Should execute with provided props',
+    );
+  });
+
   it('Shallow hierarchy (no children)', async () => {
     const doAsyncSideEffect = sinon.spy(async () => {});
     const prepareUsingProps = sinon.spy(async ({ text }) => {

src/utils/getElementType.ts

@@ -11,6 +11,7 @@ export enum ELEMENT_TYPE {
   MEMO = 7,
   FUNCTION_COMPONENT = 8,
   CLASS_COMPONENT = 9,
+  LAZY = 10,
 }
 
 function isTextNode(element: ReactNode): element is string | number {
@@ -41,6 +42,8 @@ export default function getElementType(element: ReactNode): ELEMENT_TYPE {
       return ELEMENT_TYPE.FORWARD_REF;
     } else if (type.$$typeof.toString() === 'Symbol(react.memo)') {
       return ELEMENT_TYPE.MEMO;
+    } else if (type.$$typeof.toString() === 'Symbol(react.lazy)') {
+      return ELEMENT_TYPE.LAZY;
     }
   } else if (typeof element.type === 'function') {
     if (!element.type.prototype || !('render' in element.type.prototype)) {

src/utils/reactInternalTypes.d.ts

@@ -6,6 +6,7 @@ import React, {
   ExoticComponent,
   ForwardedRef,
   ForwardRefRenderFunction,
+  LazyExoticComponent,
   MemoExoticComponent,
   Provider,
   ProviderProps,
@@ -87,3 +88,16 @@ export type MemoElement<P = unknown> = ReactElement<
   P,
   MemoExoticComponent<ComponentType<P>>
 >;
+
+type LazyPayload<P> = {
+  _status: number;
+  _result: { default: ComponentType<P> };
+};
+
+export type LazyElement<P = unknown> = ReactElement<
+  P,
+  LazyExoticComponent<ComponentType<P>> & {
+    _payload: LazyPayload<P>;
+    _init: (payload: LazyPayload<P>) => unknown;
+  }
+>;