From b21b05795fceabd0232f7d0a83b3ed2883a93f1b Mon Sep 17 00:00:00 2001
From: Swain Molster <swain.molster@lifeomic.com>
Date: Mon, 21 Aug 2023 11:14:18 -0400
Subject: [PATCH] fix: correctly support select(...) option for useAPIQuery

---
 src/hooks.test.tsx | 24 ++++++++++++++++++++++++
 src/types.ts       | 30 +++++++++++++++++++++---------
 2 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/src/hooks.test.tsx b/src/hooks.test.tsx
index 87fa8b0..07f9a13 100644
--- a/src/hooks.test.tsx
+++ b/src/hooks.test.tsx
@@ -119,6 +119,30 @@ describe('useAPIQuery', () => {
       );
     });
   });
+
+  test('using select(...) works and is typed correctly', async () => {
+    network.mock('GET /items', {
+      status: 200,
+      data: { message: 'test-message' },
+    });
+
+    const screen = render(() => {
+      const query = useAPIQuery(
+        'GET /items',
+        { filter: 'test-filter' },
+        { select: (data) => data.message },
+      );
+
+      // This line implicitly asserts that `query.data` is typed as string.
+      query.data?.codePointAt(0);
+
+      return <div data-testid="content">{query.data || ''}</div>;
+    });
+
+    await TestingLibrary.waitFor(() => {
+      expect(screen.queryByText('test-message')).toBeDefined();
+    });
+  });
 });
 
 describe('useInfiniteAPIQuery', () => {
diff --git a/src/types.ts b/src/types.ts
index 92dc8ff..2a6eb6e 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -52,15 +52,20 @@ export type RequestPayloadOf<
   Route extends keyof Endpoints,
 > = Endpoints[Route]['Request'] & PathParamsOf<Route>;
 
-type RestrictedUseQueryOptions<Response> = Omit<
-  UseQueryOptions<Response, unknown>,
-  'queryKey' | 'queryFn'
-> & {
+type RestrictedUseQueryOptions<
+  Response,
+  TError = unknown,
+  Data = Response,
+> = Omit<UseQueryOptions<Response, TError, Data>, 'queryKey' | 'queryFn'> & {
   axios?: AxiosRequestConfig;
 };
 
-type RestrictedUseInfiniteQueryOptions<Response, Request> = Omit<
-  UseInfiniteQueryOptions<Response, unknown>,
+type RestrictedUseInfiniteQueryOptions<
+  Response,
+  Request,
+  Data = Response,
+> = Omit<
+  UseInfiniteQueryOptions<Response, unknown, Data>,
   'queryKey' | 'queryFn' | 'getNextPageParam' | 'getPreviousPageParam'
 > & {
   axios?: AxiosRequestConfig;
@@ -132,11 +137,18 @@ export type CacheUtils<Endpoints extends RoughEndpoints> = {
 };
 
 export type APIQueryHooks<Endpoints extends RoughEndpoints> = {
-  useAPIQuery: <Route extends keyof Endpoints & string>(
+  useAPIQuery: <
+    Route extends keyof Endpoints & string,
+    Data = Endpoints[Route]['Response'],
+  >(
     route: Route,
     payload: RequestPayloadOf<Endpoints, Route>,
-    options?: RestrictedUseQueryOptions<Endpoints[Route]['Response']>,
-  ) => UseQueryResult<Endpoints[Route]['Response']>;
+    options?: RestrictedUseQueryOptions<
+      Endpoints[Route]['Response'],
+      unknown,
+      Data
+    >,
+  ) => UseQueryResult<Data>;
 
   useInfiniteAPIQuery: <Route extends keyof Endpoints & string>(
     route: Route,