Fix types and add tests

src/hooks/use-merged-infinite-queries/__tests__/use-merged-infinite-queries.test.ts

@@ -0,0 +1,194 @@
+import { act } from 'react-dom/test-utils';
+
+import { renderHook, waitFor } from '@/test-utils/rtl';
+
+import useMergedInfiniteQueries from '../use-merged-infinite-queries';
+import { type SingleInfiniteQueryOptions } from '../use-merged-infinite-queries.types';
+
+type MockAPIResponse = {
+  entries: Array<number>;
+  nextPage: number;
+};
+
+const MOCK_QUERY_CONFIG: Array<
+  SingleInfiniteQueryOptions<MockAPIResponse, number>
+> = [
+  {
+    queryKey: ['even-numbers'],
+    queryFn: async ({ pageParam }) => ({
+      entries: Array.from({ length: 5 }, (_, i) => pageParam + i * 2),
+      nextPage: pageParam + 10,
+    }),
+    getNextPageParam: (lastPage) => lastPage.nextPage,
+    initialPageParam: 0,
+  },
+  {
+    queryKey: ['odd-numbers'],
+    queryFn: async ({ pageParam }) => ({
+      entries: Array.from({ length: 5 }, (_, i) => pageParam + i * 2),
+      nextPage: pageParam + 10,
+    }),
+    getNextPageParam: (lastPage) => lastPage.nextPage,
+    initialPageParam: 1,
+  },
+];
+
+const MOCK_QUERY_CONFIG_WITH_ERROR: Array<
+  SingleInfiniteQueryOptions<MockAPIResponse, number>
+> = [
+  {
+    queryKey: ['even-numbers'],
+    queryFn: async ({ pageParam }) => ({
+      entries: Array.from({ length: 5 }, (_, i) => pageParam + i * 2),
+      nextPage: pageParam + 10,
+    }),
+    getNextPageParam: (lastPage) => lastPage.nextPage,
+    initialPageParam: 0,
+  },
+  {
+    queryKey: ['odd-numbers'],
+    queryFn: async () => {
+      throw new Error(`That's odd, something went wrong`);
+    },
+    getNextPageParam: (lastPage) => lastPage.nextPage,
+    initialPageParam: 1,
+  },
+];
+
+describe(useMergedInfiniteQueries.name, () => {
+  it('should merge infinite query results, and return correct loading states', async () => {
+    const { result } = renderHook(() =>
+      useMergedInfiniteQueries({
+        queries: MOCK_QUERY_CONFIG,
+        pageSize: 5,
+        flattenResponse: (res) => res.entries,
+        compare: (a, b) => (a < b ? -1 : 1),
+      })
+    );
+
+    expect(result.current[0].isFetching).toStrictEqual(true);
+    expect(result.current[0].isLoading).toStrictEqual(true);
+    expect(result.current[0].isFetchingNextPage).toStrictEqual(false);
+
+    await waitFor(() => {
+      const [mergedResult] = result.current;
+      expect(mergedResult.data).toStrictEqual([0, 1, 2, 3, 4]);
+    });
+
+    const [{ fetchNextPage }] = result.current;
+
+    act(() => {
+      fetchNextPage();
+    });
+
+    expect(result.current[0].isFetching).toStrictEqual(true);
+    expect(result.current[0].isLoading).toStrictEqual(false);
+    expect(result.current[0].isFetchingNextPage).toStrictEqual(true);
+
+    await waitFor(() => {
+      const [mergedResult] = result.current;
+      expect(mergedResult.data).toStrictEqual([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
+    });
+  });
+
+  it('should return partial query results if one of the queries fails', async () => {
+    const { result } = renderHook(() =>
+      useMergedInfiniteQueries({
+        queries: MOCK_QUERY_CONFIG_WITH_ERROR,
+        pageSize: 5,
+        flattenResponse: (res) => res.entries,
+        compare: (a, b) => (a < b ? -1 : 1),
+      })
+    );
+
+    await waitFor(() => {
+      const [mergedResult] = result.current;
+      expect(mergedResult.data).toStrictEqual([0, 2, 4, 6, 8]);
+      expect(mergedResult.status).toStrictEqual('error');
+    });
+  });
+
+  it('should fetch more data when fetchNextPage is called', async () => {
+    const { result } = renderHook(() =>
+      useMergedInfiniteQueries({
+        queries: MOCK_QUERY_CONFIG,
+        pageSize: 5,
+        flattenResponse: (res) => res.entries,
+        compare: (a, b) => (a < b ? -1 : 1),
+      })
+    );
+
+    await waitFor(() => {
+      const [_, queryResults] = result.current;
+      expect(queryResults[0].data?.pages[0].entries).toStrictEqual([
+        0, 2, 4, 6, 8,
+      ]);
+      expect(queryResults[1].data?.pages[0].entries).toStrictEqual([
+        1, 3, 5, 7, 9,
+      ]);
+    });
+
+    const [{ fetchNextPage }] = result.current;
+
+    act(() => {
+      fetchNextPage();
+    });
+
+    await waitFor(() => {
+      const [_, queryResults] = result.current;
+      expect(queryResults[0].data?.pages[1].entries).toStrictEqual([
+        10, 12, 14, 16, 18,
+      ]);
+      expect(queryResults[1].data?.pages[1].entries).toStrictEqual([
+        11, 13, 15, 17, 19,
+      ]);
+    });
+  });
+
+  it('should skip fetching more data if enough data has already been fetched', async () => {
+    const { result } = renderHook(() =>
+      useMergedInfiniteQueries({
+        queries: MOCK_QUERY_CONFIG,
+        pageSize: 5,
+        flattenResponse: (res) => res.entries,
+        compare: (a, b) => (a < b ? -1 : 1),
+      })
+    );
+
+    await waitFor(() => {
+      const [mergedResult, queryResults] = result.current;
+      expect(mergedResult.status).toBe('success');
+      // 1 page fetched from each endpoint
+      expect(queryResults[0].data?.pages.length).toStrictEqual(1);
+      expect(queryResults[1].data?.pages.length).toStrictEqual(1);
+    });
+
+    const [{ fetchNextPage: fetchSecondPage }] = result.current;
+
+    act(() => {
+      fetchSecondPage();
+    });
+
+    await waitFor(() => {
+      const [mergedResult, queryResults] = result.current;
+      expect(mergedResult.status).toBe('success');
+      // 2 pages fetched from each endpoint
+      expect(queryResults[0].data?.pages.length).toStrictEqual(2);
+      expect(queryResults[1].data?.pages.length).toStrictEqual(2);
+    });
+
+    const [{ fetchNextPage: fetchThirdPage }] = result.current;
+
+    act(() => {
+      fetchThirdPage();
+    });
+
+    await waitFor(() => {
+      const [mergedResult, queryResults] = result.current;
+      expect(mergedResult.status).toBe('success');
+      // still just 2 pages fetched from each endpoint
+      expect(queryResults[0].data?.pages.length).toStrictEqual(2);
+      expect(queryResults[1].data?.pages.length).toStrictEqual(2);
+    });
+  });
+});

src/hooks/use-merged-infinite-queries/use-merged-infinite-queries.ts

@@ -1,58 +1,52 @@
 import { useMemo, useState, useEffect } from 'react';
 
-import {
-  InfiniteQueryObserver,
-  useQueryClient,
-  type UseInfiniteQueryOptions,
-  type InfiniteQueryObserverResult,
-  type InfiniteData,
-} from '@tanstack/react-query';
+import { InfiniteQueryObserver, useQueryClient } from '@tanstack/react-query';
 
 import mergeSortedArrays from '@/utils/merge-sorted-arrays';
 
 import getMergedFetchNextPage from './helpers/get-merged-fetch-next-page';
 import getMergedQueryStatus from './helpers/get-merged-query-status';
-import { type MergedQueriesResults } from './use-merged-infinite-queries.types';
+import {
+  type SingleInfiniteQueryResult,
+  type MergedQueriesResults,
+  type Props,
+} from './use-merged-infinite-queries.types';
 
 /**
  * Combines results from multiple infinite queries in sorted order.
  *
  * To ensure the order of the combined results remains consistent even after fetching additional data,
- * this hook requests extra data from all queries. This ensures no higher-priority items are missed
+ * this hook requests extra data (page size * number of queries) from all queries. This ensures no higher-priority items are missed
  * when sorting, even if they appear in later pages of data.
  *
- * **Example:**
- * If you want to display `x` items and have `n` queries, this hook fetches `x * n` items.
  * For instance, to show 5 items across 3 queries, it fetches 15 items.
  *
  * @param queries - Array of react-query Infinite Query configurations.
  * @param pageSize - Number of items to add to the result array when requesting new data.
- * Ensure all queries return at least `pageSize` items if sufficient data is available.
  * @param flattenResult - A function that takes the expected query result and flattens it into an array of items
  * @param compare - A comparison function used to sort and merge results.
  * The function should accept two arguments and return:
- *   - A number > 0 if the second argument has a higher priority.
- *   - A number <= 0 if the first argument has a higher or equal priority.
- * **Note:** The comparison logic must match the sorting logic used in the queries to maintain consistency.
+ *   - A number > 0 if the first argument has a higher priority.
+ *   - A number <= 0 if the second argument has a higher or equal priority.
+ *   - **Note:** The comparison logic must match the sorting logic used in the queries to maintain consistency.
  *
  * @returns A tuple [mergedQueryResults, queryResults]:
  *   - `mergedQueryResults`: The merged and sorted results from all queries.
  *   - `queryResults`: An array containing individual results from each query.
  */
-export default function useMergedInfiniteQueries<T, R>({
+export default function useMergedInfiniteQueries<TData, TResponse, TPageParam>({
   queries,
   pageSize,
-  flattenResult,
+  flattenResponse,
   compare,
-}: {
-  queries: UseInfiniteQueryOptions<InfiniteData<R>>[];
-  pageSize: number;
-  flattenResult: (queryResult: R) => Array<T>;
-  compare: (a: T, b: T) => number;
-}): [MergedQueriesResults<T>, InfiniteQueryObserverResult<InfiniteData<R>>[]] {
+}: Props<TData, TResponse, TPageParam>): [
+  MergedQueriesResults<TData>,
+  Array<SingleInfiniteQueryResult<TResponse>>,
+] {
   const [count, setCount] = useState(pageSize);
+
   const [queryResults, setQueryResults] = useState<
-    Array<InfiniteQueryObserverResult<InfiniteData<R>>>
+    Array<SingleInfiniteQueryResult<TResponse>>
   >([]);
   const queryClient = useQueryClient();
 
@@ -76,15 +70,15 @@ export default function useMergedInfiniteQueries<T, R>({
     return () => unsubscribes.forEach((unsubscribe) => unsubscribe());
   }, [queries, queryClient, pageSize]);
 
-  const flattenedDataArrays: Array<Array<T>> = useMemo(() => {
+  const flattenedDataArrays: Array<Array<TData>> = useMemo(() => {
     return queryResults.map((queryResult) => {
       if (!queryResult.data) return [];
-      return queryResult.data.pages.flatMap((res) => flattenResult(res));
+      return queryResult.data.pages.flatMap((res) => flattenResponse(res));
     });
-  }, [queryResults, flattenResult]);
+  }, [queryResults, flattenResponse]);
 
   const { pointers, sortedArray } = useMemo(() => {
-    return mergeSortedArrays<T>({
+    return mergeSortedArrays<TData>({
       sortedArrays: flattenedDataArrays,
       itemsCount: count,
       compareFunc: compare,

src/hooks/use-merged-infinite-queries/use-merged-infinite-queries.types.ts

@@ -1,11 +1,39 @@
+import {
+  type QueryKey,
+  type InfiniteData,
+  type UseInfiniteQueryOptions,
+  type useInfiniteQuery,
+} from '@tanstack/react-query';
+
 export type MergedQueryStatus = 'idle' | 'loading' | 'success' | 'error';
 
-export type MergedQueriesResults<T> = {
-  data: Array<T>;
+export type MergedQueriesResults<TData> = {
+  data: Array<TData>;
   status: MergedQueryStatus;
   isLoading: boolean;
   isFetching: boolean;
   isFetchingNextPage: boolean;
   hasNextPage: boolean;
   fetchNextPage: () => void;
 };
+
+export type SingleInfiniteQueryOptions<TResponse, TPageParam> =
+  UseInfiniteQueryOptions<
+    TResponse,
+    Error,
+    InfiniteData<TResponse, TPageParam>,
+    TResponse,
+    QueryKey,
+    TPageParam
+  >;
+
+export type SingleInfiniteQueryResult<TResponse> = ReturnType<
+  typeof useInfiniteQuery<TResponse>
+>;
+
+export type Props<TData, TResponse, TPageParam> = {
+  queries: Array<SingleInfiniteQueryOptions<TResponse, TPageParam>>;
+  pageSize: number;
+  flattenResponse: (queryResult: TResponse) => Array<TData>;
+  compare: (a: TData, b: TData) => number;
+};