From 0cdf5e88b7e38239be5ff277c1094b353ed3e808 Mon Sep 17 00:00:00 2001
From: Matt Sutkowski <msutkowski@gmail.com>
Date: Sat, 6 Nov 2021 10:44:59 -0700
Subject: [PATCH] Fix types for lazy query :fingers_crossed:

---
 .../toolkit/src/query/core/buildInitiate.ts   | 27 +++++++++++++++++--
 .../toolkit/src/query/react/buildHooks.ts     |  9 ++++---
 .../src/query/tests/buildHooks.test.tsx       |  3 ---
 3 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/packages/toolkit/src/query/core/buildInitiate.ts b/packages/toolkit/src/query/core/buildInitiate.ts
index 31c5cc7e6b..85bc8ab1d1 100644
--- a/packages/toolkit/src/query/core/buildInitiate.ts
+++ b/packages/toolkit/src/query/core/buildInitiate.ts
@@ -45,9 +45,9 @@ type StartQueryActionCreator<
   options?: StartQueryActionCreatorOptions
 ) => ThunkAction<QueryActionCreatorResult<D>, any, any, AnyAction>
 
-export type QueryActionCreatorResult<
+type CommonQueryActionCreatorResult<
   D extends QueryDefinition<any, any, any, any>
-> = Promise<QuerySubState<D>> & {
+> = {
   arg: QueryArgFrom<D>
   requestId: string
   subscriptionOptions: SubscriptionOptions | undefined
@@ -58,6 +58,29 @@ export type QueryActionCreatorResult<
   updateSubscriptionOptions(options: SubscriptionOptions): void
 }
 
+export type QueryActionCreatorResult<
+  D extends QueryDefinition<any, any, any, any>
+> = Promise<QuerySubState<D>> & CommonQueryActionCreatorResult<D>
+
+export type LazyQueryActionCreatorResult<
+  D extends QueryDefinition<any, any, any, any>
+> = Promise<
+  | { data: ResultTypeFrom<D> }
+  | {
+      error:
+        | Exclude<
+            BaseQueryError<
+              D extends QueryDefinition<any, infer BaseQuery, any, any>
+                ? BaseQuery
+                : never
+            >,
+            undefined
+          >
+        | SerializedError
+    }
+> &
+  CommonQueryActionCreatorResult<D>
+
 type StartMutationActionCreator<
   D extends MutationDefinition<any, any, any, any>
 > = (
diff --git a/packages/toolkit/src/query/react/buildHooks.ts b/packages/toolkit/src/query/react/buildHooks.ts
index dd4a7c8892..bcb6a11be3 100644
--- a/packages/toolkit/src/query/react/buildHooks.ts
+++ b/packages/toolkit/src/query/react/buildHooks.ts
@@ -31,6 +31,7 @@ import type {
 import type {
   QueryActionCreatorResult,
   MutationActionCreatorResult,
+  LazyQueryActionCreatorResult,
 } from '@reduxjs/toolkit/dist/query/core/buildInitiate'
 import type { SerializeQueryArgs } from '@reduxjs/toolkit/dist/query/defaultSerializeQueryArgs'
 import { shallowEqual } from 'react-redux'
@@ -220,7 +221,7 @@ export type LazyQueryTrigger<D extends QueryDefinition<any, any, any, any>> = {
   (
     arg: QueryArgFrom<D>,
     preferCacheValue?: boolean
-  ): QueryActionCreatorResult<D>
+  ): LazyQueryActionCreatorResult<D>
 }
 
 /**
@@ -670,7 +671,7 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
       const dispatch = useDispatch<ThunkDispatch<any, any, AnyAction>>()
 
       const [arg, setArg] = useState<any>(UNINITIALIZED_VALUE)
-      const promiseRef = useRef<QueryActionCreatorResult<any> | undefined>()
+      const promiseRef = useRef<LazyQueryActionCreatorResult<any> | undefined>()
 
       const stableSubscriptionOptions = useShallowStableValue({
         refetchOnReconnect,
@@ -695,7 +696,7 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
 
       const trigger = useCallback(
         function (arg: any, preferCacheValue = false) {
-          let promise: QueryActionCreatorResult<any>
+          let promise: LazyQueryActionCreatorResult<any>
 
           batch(() => {
             promiseRef.current?.unsubscribe()
@@ -705,7 +706,7 @@ export function buildHooks<Definitions extends EndpointDefinitions>({
                 subscriptionOptions: subscriptionOptionsRef.current,
                 forceRefetch: !preferCacheValue,
               })
-            )
+            ) as any
 
             setArg(arg)
           })
diff --git a/packages/toolkit/src/query/tests/buildHooks.test.tsx b/packages/toolkit/src/query/tests/buildHooks.test.tsx
index fc2c455f17..14c5bf70ed 100644
--- a/packages/toolkit/src/query/tests/buildHooks.test.tsx
+++ b/packages/toolkit/src/query/tests/buildHooks.test.tsx
@@ -791,9 +791,6 @@ describe('hooks tests', () => {
               | {
                   error: { status: number; data: unknown } | SerializedError
                 }
-              | {
-                  data: undefined
-                }
               | {
                   data: {
                     name: string