Fix useQuery keeps polling bug from #3272

packages/hooks/src/__tests__/useQuery.test.tsx

@@ -1,7 +1,7 @@
 import React, { useState } from 'react';
 import { DocumentNode, GraphQLError } from 'graphql';
 import gql from 'graphql-tag';
-import { MockedProvider } from '@apollo/react-testing';
+import { MockedProvider, MockLink } from '@apollo/react-testing';
 import { render, cleanup } from '@testing-library/react';
 import { useQuery } from '@apollo/react-hooks';
 
@@ -143,6 +143,51 @@ describe('useQuery Hook', () => {
       );
     });
 
+    it('should stop polling when the component is unmounted', done => {
+      const mockLink = new MockLink(CAR_MOCKS);
+      const linkRequestSpy = jest.spyOn(mockLink, 'request');
+      let renderCount = 0;
+      const QueryComponent = ({ unmount }: { unmount: () => void }) => {
+        const { data, loading } = useQuery(CAR_QUERY, { pollInterval: 10 });
+        switch (renderCount) {
+          case 0:
+            expect(loading).toBeTruthy();
+            break;
+          case 1:
+            expect(loading).toBeFalsy();
+            expect(data).toEqual(CAR_RESULT_DATA);
+            expect(linkRequestSpy).toHaveBeenCalledTimes(1);
+            break;
+          case 2:
+            expect(loading).toBeFalsy();
+            expect(data).toEqual(CAR_RESULT_DATA);
+            expect(linkRequestSpy).toHaveBeenCalledTimes(2);
+            unmount();
+            break;
+          default:
+        }
+        renderCount += 1;
+        return null;
+      };
+
+      const Component = () => {
+        const [queryMounted, setQueryMounted] = useState(true);
+        const unmount = () => setTimeout(() => setQueryMounted(false), 0);
+        if (!queryMounted)
+          setTimeout(() => {
+            expect(linkRequestSpy).toHaveBeenCalledTimes(2);
+            done();
+          }, 30);
+        return <>{queryMounted && <QueryComponent unmount={unmount} />}</>;
+      };
+
+      render(
+        <MockedProvider mocks={CAR_MOCKS} link={mockLink}>
+          <Component />
+        </MockedProvider>
+      );
+    });
+
     it('should set called to true by default', () => {
       const Component = () => {
         const { loading, called } = useQuery(CAR_QUERY);

packages/hooks/src/data/QueryData.ts

@@ -112,6 +112,15 @@ export class QueryData<TData, TVariables> extends OperationData {
     this.isMounted = true;
     if (!lazy || this.runLazy) {
       this.handleErrorOrCompleted();
+
+      // When the component is done rendering stored query errors, we'll
+      // remove those errors from the `ObservableQuery` query store, so they
+      // aren't re-displayed on subsequent (potentially error free)
+      // requests/responses.
+      setTimeout(() => {
+        this.currentObservable.query &&
+          this.currentObservable.query.resetQueryStoreErrors();
+      });
     }
     return this.unmount.bind(this);
   }
@@ -383,14 +392,6 @@ export class QueryData<TData, TVariables> extends OperationData {
       }
     }
 
-    // When the component is done rendering stored query errors, we'll
-    // remove those errors from the `ObservableQuery` query store, so they
-    // aren't re-displayed on subsequent (potentially error free)
-    // requests/responses.
-    setTimeout(() => {
-      this.currentObservable.query!.resetQueryStoreErrors();
-    });
-
     result.client = this.client;
     this.previousData.loading =
       (this.previousData.result && this.previousData.result.loading) || false;

packages/hooks/src/utils/useBaseQuery.ts

@@ -41,5 +41,9 @@ export function useBaseQuery<TData = any, TVariables = OperationVariables>(
 
   useEffect(() => queryData.afterExecute({ lazy }), [result]);
 
+  useEffect(() => {
+    return () => queryData.cleanup();
+  }, []);
+
   return result;
 }

packages/testing/src/mocks/MockedProvider.tsx

@@ -3,6 +3,7 @@ import { ApolloClient, DefaultOptions, Resolvers } from 'apollo-client';
 import { ApolloCache } from 'apollo-cache';
 import { InMemoryCache as Cache } from 'apollo-cache-inmemory';
 import { ApolloProvider } from '@apollo/react-common';
+import { ApolloLink } from 'apollo-link';
 import { MockLink } from './mockLink';
 import { MockedResponse } from './types';
 
@@ -14,6 +15,7 @@ export interface MockedProviderProps<TSerializedCache = {}> {
   resolvers?: Resolvers;
   childProps?: object;
   children?: React.ReactElement;
+  link?: ApolloLink;
 }
 
 export interface MockedProviderState {
@@ -23,19 +25,26 @@ export interface MockedProviderState {
 export class MockedProvider extends React.Component<
   MockedProviderProps,
   MockedProviderState
-  > {
+> {
   public static defaultProps: MockedProviderProps = {
     addTypename: true
   };
 
   constructor(props: MockedProviderProps) {
     super(props);
 
-    const { mocks, addTypename, defaultOptions, cache, resolvers } = this.props;
+    const {
+      mocks,
+      addTypename,
+      defaultOptions,
+      cache,
+      resolvers,
+      link
+    } = this.props;
     const client = new ApolloClient({
       cache: cache || new Cache({ addTypename }),
       defaultOptions,
-      link: new MockLink(mocks || [], addTypename),
+      link: link || new MockLink(mocks || [], addTypename),
       resolvers
     });