From e8ee9393093d9cfcd564d7b10043243b595d3896 Mon Sep 17 00:00:00 2001
From: Brian Kim <briankimpossible@gmail.com>
Date: Thu, 30 Sep 2021 12:49:15 -0400
Subject: [PATCH] Make useLazyQuery execution function return a promise

Implements #7714
---
 .../hooks/__tests__/useLazyQuery.test.tsx     | 38 +++++++++++++++++++
 src/react/hooks/useLazyQuery.ts               | 36 +++++++++++++++---
 src/react/types/types.ts                      |  2 +-
 3 files changed, 69 insertions(+), 7 deletions(-)

diff --git a/src/react/hooks/__tests__/useLazyQuery.test.tsx b/src/react/hooks/__tests__/useLazyQuery.test.tsx
index 3ced64c9cbd..880d96bff40 100644
--- a/src/react/hooks/__tests__/useLazyQuery.test.tsx
+++ b/src/react/hooks/__tests__/useLazyQuery.test.tsx
@@ -506,4 +506,42 @@ describe('useLazyQuery Hook', () => {
     expect(result.current[1].loading).toBe(false);
     expect(result.current[1].data).toEqual({ hello: 'from link' });
   });
+
+  it('should return a promise from the execution function which resolves with the result', async () => {
+    const query = gql`{ hello }`;
+    const mocks = [
+      {
+        request: { query },
+        result: { data: { hello: 'world' } },
+        delay: 20,
+      },
+    ];
+    const { result, waitForNextUpdate } = renderHook(
+      () => useLazyQuery(query),
+      {
+        wrapper: ({ children }) => (
+          <MockedProvider mocks={mocks}>
+            {children}
+          </MockedProvider>
+        ),
+      },
+    );
+
+    expect(result.current[1].loading).toBe(false);
+    expect(result.current[1].data).toBe(undefined);
+    const execute = result.current[0];
+    const mock = jest.fn();
+    setTimeout(() => mock(execute()));
+
+    await waitForNextUpdate();
+    expect(result.current[1].loading).toBe(true);
+
+    await waitForNextUpdate();
+    expect(result.current[1].loading).toBe(false);
+    expect(result.current[1].data).toEqual({ hello: 'world' });
+
+    expect(mock).toHaveBeenCalledTimes(1);
+    expect(mock.mock.calls[0][0]).toBeInstanceOf(Promise);
+    expect(await mock.mock.calls[0][0]).toEqual(result.current[1]);
+  });
 });
diff --git a/src/react/hooks/useLazyQuery.ts b/src/react/hooks/useLazyQuery.ts
index c6db275f014..732132621dc 100644
--- a/src/react/hooks/useLazyQuery.ts
+++ b/src/react/hooks/useLazyQuery.ts
@@ -1,6 +1,6 @@
 import { DocumentNode } from 'graphql';
 import { TypedDocumentNode } from '@graphql-typed-document-node/core';
-import { useCallback, useState } from 'react';
+import { useCallback, useEffect, useState } from 'react';
 
 import {
   LazyQueryHookOptions,
@@ -26,31 +26,55 @@ export function useLazyQuery<TData = any, TVariables = OperationVariables>(
   options?: LazyQueryHookOptions<TData, TVariables>
 ): QueryTuple<TData, TVariables> {
   const [execution, setExecution] = useState<
-    { called: boolean, options?: QueryLazyOptions<TVariables> }
+    {
+      called: boolean,
+      options?: QueryLazyOptions<TVariables>,
+      resolves: Array<(result: LazyQueryResult<TData, TVariables>) => void>,
+    }
   >({
     called: false,
+    resolves: [],
   });
 
   const execute = useCallback<
     QueryTuple<TData, TVariables>[0]
   >((executeOptions?: QueryLazyOptions<TVariables>) => {
+    let resolve!: (result: LazyQueryResult<TData, TVariables>) => void;
+    const promise = new Promise<LazyQueryResult<TData, TVariables>>(
+      (resolve1) => (resolve = resolve1),
+    );
     setExecution((execution) => {
       if (execution.called) {
         result && result.refetch(executeOptions?.variables);
       }
 
-      return { called: true, options: executeOptions };
+      return {
+        called: true,
+        resolves: [...execution.resolves, resolve],
+        options: executeOptions,
+      };
     });
+
+    return promise;
   }, []);
 
-  let result = useQuery<TData, TVariables>(query, {
+  const queryOptions = {
     ...options,
     ...execution.options,
     // We don’t set skip to execution.called, because we need useQuery to call
     // addQueryPromise, so that ssr calls waits for execute to be called.
     fetchPolicy: execution.called ? options?.fetchPolicy : 'standby',
     skip: undefined,
-  });
+  };
+
+  let result = useQuery<TData, TVariables>(query, queryOptions);
+  useEffect(() => {
+    const { resolves } = execution;
+    if (!result.loading && resolves.length) {
+      setExecution((execution) => ({ ...execution, resolves: [] }));
+      resolves.forEach((resolve) => resolve(result));
+    }
+  }, [result, execution]);
 
   if (!execution.called) {
     result = {
@@ -66,7 +90,7 @@ export function useLazyQuery<TData = any, TVariables = OperationVariables>(
     for (const key of EAGER_METHODS) {
       const method = result[key];
       result[key] = (...args: any) => {
-        setExecution({ called: true });
+        setExecution((execution) => ({ ...execution, called: true }));
         return (method as any)(...args);
       };
     }
diff --git a/src/react/types/types.ts b/src/react/types/types.ts
index 36ba1993dc8..6ba2db9919e 100644
--- a/src/react/types/types.ts
+++ b/src/react/types/types.ts
@@ -108,7 +108,7 @@ export interface QueryLazyOptions<TVariables> {
 export type LazyQueryResult<TData, TVariables> = QueryResult<TData, TVariables>;
 
 export type QueryTuple<TData, TVariables> = [
-  (options?: QueryLazyOptions<TVariables>) => void,
+  (options?: QueryLazyOptions<TVariables>) => Promise<LazyQueryResult<TData, TVariables>>,
   LazyQueryResult<TData, TVariables>
 ];