diff --git a/packages/hooks/src/__tests__/useLazyQuery.test.tsx b/packages/hooks/src/__tests__/useLazyQuery.test.tsx index f9a1f39945..af28a9dbf0 100644 --- a/packages/hooks/src/__tests__/useLazyQuery.test.tsx +++ b/packages/hooks/src/__tests__/useLazyQuery.test.tsx @@ -447,4 +447,65 @@ describe('useLazyQuery Hook', () => { expect(renderCount).toBe(5); }); }); + + it('should pass updateQuery in QueryResult', async () => { + const data1 = CAR_RESULT_DATA; + const data2 = { + cars: [] + }; + const mocks = [ + { + request: { + query: CAR_QUERY + }, + result: { data: data1 } + } + ]; + + let renderCount = 0; + const Component = () => { + const [execute, { loading, data, updateQuery }] = useLazyQuery( + CAR_QUERY, + { + fetchPolicy: 'network-only' + } + ); + switch (renderCount) { + case 0: + expect(updateQuery).toBeDefined(); + expect(loading).toEqual(false); + setTimeout(() => { + execute(); + }); + break; + case 1: + expect(loading).toEqual(true); + break; + case 2: + expect(loading).toEqual(false); + expect(data).toEqual(data1); + setTimeout(() => { + updateQuery(() => data2); + }); + break; + case 3: + expect(loading).toEqual(false); + expect(data).toEqual(data2); + break; + default: // Do nothing + } + renderCount += 1; + return null; + }; + + render( + + + + ); + + await wait(() => { + expect(renderCount).toBe(4); + }); + }); }); diff --git a/packages/hooks/src/data/QueryData.ts b/packages/hooks/src/data/QueryData.ts index 1cb88001e2..426a1ab516 100644 --- a/packages/hooks/src/data/QueryData.ts +++ b/packages/hooks/src/data/QueryData.ts @@ -14,6 +14,7 @@ import { QueryResult, ObservableQueryFields } from '@apollo/react-common'; +import { invariant } from 'ts-invariant'; import { QueryPreviousData, @@ -69,7 +70,8 @@ export class QueryData extends OperationData { loading: false, networkStatus: NetworkStatus.ready, called: false, - data: undefined + data: undefined, + updateQuery: this.lazyUpdateQuery } as QueryResult ] : [this.runLazyQuery, this.execute()]; @@ -487,4 +489,16 @@ export class QueryData extends OperationData { subscribeToMore: this.obsSubscribeToMore } as ObservableQueryFields; } + + private lazyUpdateQuery: QueryResult< + TData, + TVariables + >['updateQuery'] = mapFn => { + const { query } = this.currentObservable; + invariant( + query, + `updateQuery cannot be invoked if the lazy query was not executed first.` + ); + return query!.updateQuery(mapFn); + }; }