From cf90a3d8c67a953888e984a9ee6dc0a0d7e884ff Mon Sep 17 00:00:00 2001
From: Brian Kim <briankimpossible@gmail.com>
Date: Wed, 18 Aug 2021 17:35:37 -0400
Subject: [PATCH] delay state updates until after the mutation execution
 promise resolves

Fixes #8137
---
 .../hooks/__tests__/useMutation.test.tsx      |  38 +++++++
 src/react/hooks/useMutation.ts                | 106 ++++++++++--------
 2 files changed, 96 insertions(+), 48 deletions(-)

diff --git a/src/react/hooks/__tests__/useMutation.test.tsx b/src/react/hooks/__tests__/useMutation.test.tsx
index 82c3e081367..b5e41e20561 100644
--- a/src/react/hooks/__tests__/useMutation.test.tsx
+++ b/src/react/hooks/__tests__/useMutation.test.tsx
@@ -188,6 +188,44 @@ describe('useMutation Hook', () => {
       });
     });
 
+    it('the useMutation promise should resolve before the hook updates', async () => {
+      const variables = {
+        description: 'Get milk!'
+      };
+
+      const mocks = [
+        {
+          request: {
+            query: CREATE_TODO_MUTATION,
+            variables
+          },
+          result: { data: CREATE_TODO_RESULT }
+        }
+      ];
+
+      const { result, waitForNextUpdate } = renderHook(
+        () => useMutation(CREATE_TODO_MUTATION),
+        { wrapper: ({ children }) => (
+          <MockedProvider mocks={mocks}>
+            {children}
+          </MockedProvider>
+        )},
+      );
+
+      const mutate = result.current[0];
+
+      await act(async () => {
+        await mutate({ variables });
+      });
+
+      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(CREATE_TODO_RESULT);
+    });
+
     describe('mutate function upon error', () => {
       it('resolves with the resulting data and errors', async () => {
         const variables = {
diff --git a/src/react/hooks/useMutation.ts b/src/react/hooks/useMutation.ts
index ade44580ebf..d1c25ddb910 100644
--- a/src/react/hooks/useMutation.ts
+++ b/src/react/hooks/useMutation.ts
@@ -39,6 +39,7 @@ export function useMutation<
   const ref = useRef({
     result,
     mutationId: 0,
+    baseOptions: { ...options, mutation },
     isMounted: true,
   });
 
@@ -50,7 +51,7 @@ export function useMutation<
       TCache
     > = {},
   ) => {
-    if (!ref.current.result.loading && !options?.ignoreResults) {
+    if (!ref.current.result.loading && !ref.current.baseOptions?.ignoreResults) {
       setResult(ref.current.result = {
         loading: true,
         error: void 0,
@@ -62,64 +63,73 @@ export function useMutation<
 
     const mutationId = ++ref.current.mutationId;
     const clientOptions = mergeOptions(
-      { ...options, mutation },
+      ref.current.baseOptions,
       executeOptions as any,
     );
-    return client.mutate(clientOptions).then((response) => {
-      const { data, errors } = response;
-      const error =
-        errors && errors.length > 0
-          ? new ApolloError({ graphQLErrors: errors })
-          : void 0;
 
-      if (mutationId === ref.current.mutationId && !options?.ignoreResults) {
-        const result = {
-          called: true,
-          loading: false,
-          data,
-          error,
-          client,
-        };
-
-        if (
-          ref.current.isMounted &&
-          !equal(ref.current.result, result)
-        ) {
-          ref.current.result = result;
-          setResult(result);
-        }
-      }
-
-      options?.onCompleted?.(data!);
+    const execution = client.mutate(clientOptions).then((response) =>{
+      ref.current.baseOptions?.onCompleted?.(response.data!);
       return response;
     }).catch((error) => {
-      if (mutationId === ref.current.mutationId) {
-        const result = {
-          loading: false,
-          error,
-          data: void 0,
-          called: true,
-          client,
-        };
-
-        if (
-          ref.current.isMounted &&
-          !equal(ref.current.result, result)
-        ) {
-          ref.current.result = result;
-          setResult(result);
-        }
-      }
-
-      if (options?.onError) {
-        options.onError(error);
+      if (ref.current.baseOptions?.onError) {
+        ref.current.baseOptions?.onError(error);
         // TODO(brian): why are we returning this here???
         return { data: void 0, errors: error };
       }
 
       throw error;
     });
-  }, [client, options, mutation]);
+
+    execution
+      .finally(() => new Promise((resolve) => setTimeout(resolve)))
+      .then(
+        (response) => {
+          const { data, errors } = response;
+          const error =
+            errors && errors.length > 0
+              ? new ApolloError({ graphQLErrors: errors })
+              : void 0;
+
+          if (
+            mutationId === ref.current.mutationId &&
+            !ref.current.baseOptions?.ignoreResults
+          ) {
+            const result = {
+              called: true,
+              loading: false,
+              data,
+              error,
+              client,
+            };
+
+            if (ref.current.isMounted && !equal(ref.current.result, result)) {
+              setResult(ref.current.result = result);
+            }
+          }
+        },
+        (error) => {
+          if (mutationId === ref.current.mutationId) {
+            const result = {
+              loading: false,
+              error,
+              data: void 0,
+              called: true,
+              client,
+            };
+
+            if (ref.current.isMounted && !equal(ref.current.result, result)) {
+              setResult(ref.current.result = result);
+            }
+          }
+        },
+      );
+
+    return execution;
+  }, [client]);
+
+  useEffect(() => {
+    ref.current.baseOptions = { ...options, mutation };
+  }, [options, mutation]);
 
   useEffect(() => () => {
     ref.current.isMounted = false;