From b010a1522e8a00ab83eb613e2efe467e74c58756 Mon Sep 17 00:00:00 2001
From: Paul Venable <pvenable@gmail.com>
Date: Wed, 4 May 2022 10:36:31 -0700
Subject: [PATCH] Support Query and QueryVariables type arguments in
 mockGraphQLQuery and mockGraphQLMutation (#5274)

Co-authored-by: Aditya Pandey <aditya@autospace.co>
---
 packages/testing/src/web/mockRequests.ts | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/packages/testing/src/web/mockRequests.ts b/packages/testing/src/web/mockRequests.ts
index c3036296e4d2..dfb076d1e9fd 100644
--- a/packages/testing/src/web/mockRequests.ts
+++ b/packages/testing/src/web/mockRequests.ts
@@ -67,8 +67,11 @@ export const registerHandler = (handler: RequestHandler) => {
   }
 }
 
-export type DataFunction = (
-  variables: Record<string, any>,
+export type DataFunction<
+  Query extends Record<string, unknown> = Record<string, unknown>,
+  QueryVariables = Record<string, any>
+> = (
+  variables: QueryVariables,
   {
     req,
     ctx,
@@ -76,7 +79,7 @@ export type DataFunction = (
     req: GraphQLRequest<any>
     ctx: GraphQLContext<Record<string, any>>
   }
-) => Record<string, unknown> | void
+) => Query | void
 
 // These should get exported from MSW
 type ResponseFunction<BodyType = any> = (
@@ -144,17 +147,23 @@ const mockGraphQL = (
   return data
 }
 
-export const mockGraphQLQuery = (
+export const mockGraphQLQuery = <
+  Query extends Record<string, unknown> = Record<string, unknown>,
+  QueryVariables = Record<string, any>
+>(
   operation: string,
-  data: DataFunction | Record<string, unknown>,
+  data: DataFunction<Query, QueryVariables> | Query,
   responseEnhancer?: ResponseEnhancer
 ) => {
   return mockGraphQL('query', operation, data, responseEnhancer)
 }
 
-export const mockGraphQLMutation = (
+export const mockGraphQLMutation = <
+  Query extends Record<string, unknown> = Record<string, unknown>,
+  QueryVariables = Record<string, any>
+>(
   operation: string,
-  data: DataFunction | Record<string, unknown>,
+  data: DataFunction<Query, QueryVariables> | Query,
   responseEnhancer?: ResponseEnhancer
 ) => {
   return mockGraphQL('mutation', operation, data, responseEnhancer)