From b6309e7f84524c5f35acdb42912eadf55fda9da0 Mon Sep 17 00:00:00 2001
From: Brian Kim <briankimpossible@gmail.com>
Date: Thu, 12 Aug 2021 21:26:43 -0400
Subject: [PATCH] Allow observableQuery methods to be called by useLazyQuery

Fixes #5140
---
 .../hooks/__tests__/useLazyQuery.test.tsx     | 54 +++++++++++++++++++
 src/react/hooks/useLazyQuery.ts               | 10 ++++
 src/react/types/types.ts                      | 25 +--------
 3 files changed, 66 insertions(+), 23 deletions(-)

diff --git a/src/react/hooks/__tests__/useLazyQuery.test.tsx b/src/react/hooks/__tests__/useLazyQuery.test.tsx
index c5c9652d73f..81f58cfb638 100644
--- a/src/react/hooks/__tests__/useLazyQuery.test.tsx
+++ b/src/react/hooks/__tests__/useLazyQuery.test.tsx
@@ -327,4 +327,58 @@ describe('useLazyQuery Hook', () => {
     expect(result.current[1].data).toEqual({ hello: 'world 2' });
     expect(result.current[1].previousData).toEqual({ hello: 'world 1' });
   });
+
+  it('should allow for the query to start with polling', async () => {
+    const query = gql`{ hello }`;
+    const mocks = [
+      {
+        request: { query },
+        result: { data: { hello: "world 1" } },
+      },
+      {
+        request: { query },
+        result: { data: { hello: "world 2" } },
+      },
+      {
+        request: { query },
+        result: { data: { hello: "world 3" } },
+      },
+    ];
+
+    const wrapper = ({ children }: any) => (
+      <MockedProvider mocks={mocks}>{children}</MockedProvider>
+    );
+
+    const { result, waitForNextUpdate } = renderHook(
+      () => useLazyQuery(query),
+      { wrapper },
+    );
+
+    expect(result.current[1].loading).toBe(false);
+    expect(result.current[1].data).toBe(undefined);
+
+    setTimeout(() => {
+      result.current[1].startPolling(10);
+    });
+
+    await waitForNextUpdate();
+    expect(result.current[1].loading).toBe(true);
+    expect(result.current[1].data).toBe(undefined);
+
+    await waitForNextUpdate();
+    expect(result.current[1].loading).toBe(false);
+    expect(result.current[1].data).toEqual({ hello: "world 1" });
+
+    await waitForNextUpdate();
+    expect(result.current[1].loading).toBe(false);
+    expect(result.current[1].data).toEqual({ hello: "world 2" });
+
+    await waitForNextUpdate();
+
+    expect(result.current[1].loading).toBe(false);
+    expect(result.current[1].data).toEqual({ hello: "world 3" });
+
+    result.current[1].stopPolling();
+    await expect(waitForNextUpdate({ timeout: 20 })).rejects.toThrow('Timed out');
+  });
 });
diff --git a/src/react/hooks/useLazyQuery.ts b/src/react/hooks/useLazyQuery.ts
index 9e5b76532bd..1762d11694b 100644
--- a/src/react/hooks/useLazyQuery.ts
+++ b/src/react/hooks/useLazyQuery.ts
@@ -43,6 +43,16 @@ export function useLazyQuery<TData = any, TVariables = OperationVariables>(
   });
 
   if (!execution.called) {
+    for (const key in result) {
+      if (typeof (result as any)[key] === 'function') {
+        const method = (result as any)[key];
+        (result as any)[key] = (...args: any) => {
+          setExecution({ called: true });
+          return method(...args);
+        };
+      }
+    }
+
     result = {
       ...result,
       loading: false,
diff --git a/src/react/types/types.ts b/src/react/types/types.ts
index 552784aa749..3dbd527e6a0 100644
--- a/src/react/types/types.ts
+++ b/src/react/types/types.ts
@@ -104,29 +104,8 @@ export interface QueryLazyOptions<TVariables> {
   context?: DefaultContext;
 }
 
-type UnexecutedLazyFields = {
-  loading: false;
-  networkStatus: NetworkStatus.ready;
-  called: false;
-  data: undefined;
-  previousData?: undefined;
-}
-
-type Impartial<T> = {
-  [P in keyof T]?: never;
-}
-
-type AbsentLazyResultFields =
-  Omit<
-    Impartial<QueryResult<unknown, unknown>>,
-    keyof UnexecutedLazyFields>
-
-type UnexecutedLazyResult =
-   UnexecutedLazyFields & AbsentLazyResultFields
-
-export type LazyQueryResult<TData, TVariables> =
-  | UnexecutedLazyResult
-  | QueryResult<TData, TVariables>;
+// TODO: Delete this
+export type LazyQueryResult<TData, TVariables> = QueryResult<TData, TVariables>;
 
 export type QueryTuple<TData, TVariables> = [
   (options?: QueryLazyOptions<TVariables>) => void,