Add suspense to wrapRootComponent

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ableco/abledev-react",
-  "version": "0.0.1-alpha6",
+  "version": "0.0.1-alpha7",
   "description": "",
   "main": "dist/abledev-react.js",
   "types": "dist/index.d.ts",

src/useQuery.test.tsx

@@ -33,6 +33,29 @@ const TestComponent = wrapRootComponent(() => {
   );
 });
 
+const SuspenseTestComponent = () => {
+  const { data } = useUnsweetenedQuery("queries/something");
+  return (
+    <>
+      <span>done</span>
+      <span>{JSON.stringify(data as object)}</span>
+    </>
+  );
+};
+
+const InternalSuspenseTestComponent = wrapRootComponent(SuspenseTestComponent, {
+  suspense: {
+    mode: "withInternalLoadingState",
+    loadingUI: <span>loading...</span>,
+  },
+});
+
+const ExternalSuspenseTestComponent = wrapRootComponent(SuspenseTestComponent, {
+  suspense: {
+    mode: "withExternalLoadingState",
+  },
+});
+
 const ParameterizedComponent = wrapRootComponent(
   ({ name }: { name: string }) => {
     const { isLoading, data } = useUnsweetenedQuery("queries/something", {
@@ -57,7 +80,39 @@ test("It calls a server", async () => {
   expect(await screen.findByText(/queries\/something/)).toBeInTheDocument();
 });
 
+test("It shows a loading UI if suspense with internal loading state", async () => {
+  render(<InternalSuspenseTestComponent />);
+  expect(screen.getByText("loading...")).toBeInTheDocument();
+  expect(await screen.findByText("done")).toBeInTheDocument();
+  expect(await screen.findByText(/queries\/something/)).toBeInTheDocument();
+});
+
+test("It requires to be wrapped in Suspense if suspense with external state", async () => {
+  expectToThrow(() => {
+    render(<ExternalSuspenseTestComponent />);
+  });
+
+  render(
+    <React.Suspense fallback={<span>loading...</span>}>
+      <ExternalSuspenseTestComponent />
+    </React.Suspense>,
+  );
+
+  expect(screen.getByText("loading...")).toBeInTheDocument();
+  expect(await screen.findByText("done")).toBeInTheDocument();
+  expect(await screen.findByText(/queries\/something/)).toBeInTheDocument();
+});
+
 test("It can receive some arguments", async () => {
   render(<ParameterizedComponent name="Rare name" />);
   expect(await screen.findByText(/Rare name/)).toBeInTheDocument();
 });
+
+const expectToThrow = (func: Function) => {
+  // Even though the error is caught, it still gets printed to the console
+  // so we mock that out to avoid the wall of red text.
+  const spy = jest.spyOn(console, "error");
+  spy.mockImplementation(() => {});
+  expect(func).toThrow();
+  spy.mockRestore();
+};

src/wrapRootComponent.tsx

@@ -1,14 +1,51 @@
 import * as React from "react";
 import { QueryClient, QueryClientProvider } from "react-query";
 
-function wrapRootComponent<PropsType>(Component: React.FC<PropsType>) {
+type TurnedOffSuspense = {
+  mode: "off";
+  loadingUI?: never;
+};
+
+type WithInternalLoadingStateSuspense = {
+  mode: "withInternalLoadingState";
+  loadingUI: React.ReactElement;
+};
+
+type WithExternalLoadingStateSuspense = {
+  mode: "withExternalLoadingState";
+  loadingUI?: never;
+};
+
+type Options = {
+  suspense:
+    | TurnedOffSuspense
+    | WithInternalLoadingStateSuspense
+    | WithExternalLoadingStateSuspense;
+};
+
+function wrapRootComponent<ComponentProps>(
+  Component: React.FC<ComponentProps>,
+  options: Options = { suspense: { mode: "off" } },
+) {
   let queryClient: QueryClient;
+  const { suspense } = options;
+
+  const AbledevWrapper = (props: ComponentProps) => {
+    queryClient ??= new QueryClient({
+      defaultOptions: {
+        queries: { suspense: suspense.mode !== "off" },
+      },
+    });
 
-  const AbledevWrapper = (props: PropsType) => {
-    queryClient ??= new QueryClient();
     return (
       <QueryClientProvider client={queryClient}>
-        <Component {...props} />
+        {suspense.mode === "withInternalLoadingState" ? (
+          <React.Suspense fallback={suspense.loadingUI}>
+            <Component {...props} />
+          </React.Suspense>
+        ) : (
+          <Component {...props} />
+        )}
       </QueryClientProvider>
     );
   };