From 6f8dc958fe762aaa16f8fc6d3d3a680b7ad723da Mon Sep 17 00:00:00 2001
From: Jenn Creighton <jenn@apollographql.com>
Date: Thu, 1 Apr 2021 17:06:45 -0400
Subject: [PATCH] Fix failing tests

---
 src/core/__tests__/QueryManager/index.ts |  4 ++--
 src/core/types.ts                        |  2 +-
 src/react/components/types.ts            |  9 +++++----
 src/react/hoc/mutation-hoc.tsx           | 14 ++++++++------
 src/react/hoc/types.ts                   |  8 +++++---
 5 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/src/core/__tests__/QueryManager/index.ts b/src/core/__tests__/QueryManager/index.ts
index f1276b45da0..789e4ef8342 100644
--- a/src/core/__tests__/QueryManager/index.ts
+++ b/src/core/__tests__/QueryManager/index.ts
@@ -5082,7 +5082,7 @@ describe('QueryManager', () => {
 
   describe('awaitRefetchQueries', () => {
     const awaitRefetchTest =
-    ({ awaitRefetchQueries, testQueryError = false }: MutationBaseOptions & { testQueryError?: boolean }) =>
+    ({ awaitRefetchQueries, testQueryError = false }: MutationBaseOptions<any, any, any> & { testQueryError?: boolean }) =>
     new Promise((resolve, reject) => {
       const query = gql`
         query getAuthors($id: ID!) {
@@ -5156,7 +5156,7 @@ describe('QueryManager', () => {
         { observable },
         result => {
           expect(stripSymbols(result.data)).toEqual(queryData);
-          const mutateOptions: MutationOptions = {
+          const mutateOptions: MutationOptions<any, any, any> = {
             mutation,
             refetchQueries: ['getAuthors'],
           };
diff --git a/src/core/types.ts b/src/core/types.ts
index fa85040d32f..49418ad91e4 100644
--- a/src/core/types.ts
+++ b/src/core/types.ts
@@ -47,7 +47,7 @@ export type MutationQueryReducersMap<T = { [key: string]: any }> = {
   [queryName: string]: MutationQueryReducer<T>;
 };
 
-export type MutationUpdaterFunction<T extends Record<string, any>, TVariables, TContext> = (
+export type MutationUpdaterFunction<T, TVariables, TContext> = (
   cache: ApolloCache<T>,
   result: Omit<FetchResult<T>, 'context'>,
   options: {
diff --git a/src/react/components/types.ts b/src/react/components/types.ts
index 83dea55043f..6f602afb593 100644
--- a/src/react/components/types.ts
+++ b/src/react/components/types.ts
@@ -1,7 +1,7 @@
 import { DocumentNode } from 'graphql';
 import { TypedDocumentNode } from '@graphql-typed-document-node/core';
 
-import { OperationVariables } from '../../core';
+import { OperationVariables, Context } from '../../core';
 import {
   QueryFunctionOptions,
   QueryResult,
@@ -22,11 +22,12 @@ export interface QueryComponentOptions<
 
 export interface MutationComponentOptions<
   TData = any,
-  TVariables = OperationVariables
-> extends BaseMutationOptions<TData, TVariables> {
+  TVariables = OperationVariables,
+  TContext = Context,
+> extends BaseMutationOptions<TData, TVariables, TContext> {
   mutation: DocumentNode | TypedDocumentNode<TData, TVariables>;
   children: (
-    mutateFunction: MutationFunction<TData, TVariables>,
+    mutateFunction: MutationFunction<TData, TVariables, TContext>,
     result: MutationResult<TData>
   ) => JSX.Element | null;
 }
diff --git a/src/react/hoc/mutation-hoc.tsx b/src/react/hoc/mutation-hoc.tsx
index 5159753e678..9c88fe10a23 100644
--- a/src/react/hoc/mutation-hoc.tsx
+++ b/src/react/hoc/mutation-hoc.tsx
@@ -3,6 +3,7 @@ import { DocumentNode } from 'graphql';
 import hoistNonReactStatics from 'hoist-non-react-statics';
 
 import { parser } from '../parser';
+import { Context } from '../../core/types';
 import {
   BaseMutationOptions,
   MutationFunction,
@@ -20,9 +21,10 @@ import { OperationOption, OptionProps, MutateProps } from './types';
 
 export function withMutation<
   TProps extends TGraphQLVariables | {} = {},
-  TData = {},
+  TData extends Record<string, any> = {},
   TGraphQLVariables = {},
-  TChildProps = MutateProps<TData, TGraphQLVariables>
+  TChildProps = MutateProps<TData, TGraphQLVariables>,
+  TContext = Context,
 >(
   document: DocumentNode,
   operationOptions: OperationOption<
@@ -41,9 +43,9 @@ export function withMutation<
     alias = 'Apollo'
   } = operationOptions;
 
-  let mapPropsToOptions = options as (props: any) => BaseMutationOptions;
+  let mapPropsToOptions = options as (props: any) => BaseMutationOptions<TData, TGraphQLVariables, TContext>;
   if (typeof mapPropsToOptions !== 'function')
-    mapPropsToOptions = () => options as BaseMutationOptions;
+    mapPropsToOptions = () => options as BaseMutationOptions<TData, TGraphQLVariables, TContext>;
 
   return (
     WrappedComponent: React.ComponentType<TProps & TChildProps>
@@ -54,7 +56,7 @@ export function withMutation<
       static WrappedComponent = WrappedComponent;
       render() {
         let props = this.props as TProps;
-        const opts = mapPropsToOptions(props);
+        const opts = mapPropsToOptions(props) as BaseMutationOptions<TData, TGraphQLVariables, TContext>;
 
         if (operationOptions.withRef) {
           this.withRef = true;
@@ -63,7 +65,7 @@ export function withMutation<
           });
         }
         if (!opts.variables && operation.variables.length > 0) {
-          opts.variables = calculateVariablesFromProps(operation, props);
+          opts.variables = calculateVariablesFromProps(operation, props) as TGraphQLVariables;
         }
 
         return (
diff --git a/src/react/hoc/types.ts b/src/react/hoc/types.ts
index 4c0ec8c3ec8..3b75f8a6fad 100644
--- a/src/react/hoc/types.ts
+++ b/src/react/hoc/types.ts
@@ -7,6 +7,7 @@ import {
   UpdateQueryOptions,
   FetchMoreQueryOptions,
   SubscribeToMoreOptions,
+  Context,
 } from '../../core';
 import {
   MutationFunction,
@@ -89,16 +90,17 @@ export interface OperationOption<
   TProps,
   TData,
   TGraphQLVariables = OperationVariables,
-  TChildProps = ChildProps<TProps, TData, TGraphQLVariables>
+  TChildProps = ChildProps<TProps, TData, TGraphQLVariables>,
+  TContext = Context,
 > {
   options?:
     | BaseQueryOptions<TGraphQLVariables>
-    | BaseMutationOptions<TData, TGraphQLVariables>
+    | BaseMutationOptions<TData, TGraphQLVariables, TContext>
     | ((
         props: TProps
       ) =>
         | BaseQueryOptions<TGraphQLVariables>
-        | BaseMutationOptions<TData, TGraphQLVariables>
+        | BaseMutationOptions<TData, TGraphQLVariables, TContext>
       );
   props?: (
     props: OptionProps<TProps, TData, TGraphQLVariables>,