fix(mutations): avoid infinite loading states if callbacks return an error (#3343)

* fix(mutations): avoid infinite loading states if callbacks return an error

add failing test cases

* fix(mutations): avoid infinite loading states if callbacks return an error

by making sure we always dispatch the error to go to error state internally;
re-writing to async-await because it has better support than promise.finally, and the flow is also easier to reason about here

* fix(mutations): fix merge conflicts

Author: TkDodo

src/core/mutation.ts

@@ -160,35 +160,7 @@ export class Mutation<
     return this.execute()
   }
 
-  execute(): Promise<TData> {
-    let data: TData
-
-    const restored = this.state.status === 'loading'
-
-    let promise = Promise.resolve()
-
-    if (!restored) {
-      this.dispatch({ type: 'loading', variables: this.options.variables! })
-      promise = promise
-        .then(() => {
-          // Notify cache callback
-          this.mutationCache.config.onMutate?.(
-            this.state.variables,
-            this as Mutation<unknown, unknown, unknown, unknown>
-          )
-        })
-        .then(() => this.options.onMutate?.(this.state.variables!))
-        .then(context => {
-          if (context !== this.state.context) {
-            this.dispatch({
-              type: 'loading',
-              context,
-              variables: this.state.variables,
-            })
-          }
-        })
-    }
-
+  async execute(): Promise<TData> {
     const executeMutation = () => {
       this.retryer = createRetryer({
         fn: () => {
@@ -214,38 +186,51 @@ export class Mutation<
       return this.retryer.promise
     }
 
-    return promise
-      .then(executeMutation)
-      .then(result => {
-        data = result
+    const restored = this.state.status === 'loading'
+    try {
+      if (!restored) {
+        this.dispatch({ type: 'loading', variables: this.options.variables! })
         // Notify cache callback
-        this.mutationCache.config.onSuccess?.(
-          data,
+        this.mutationCache.config.onMutate?.(
           this.state.variables,
-          this.state.context,
           this as Mutation<unknown, unknown, unknown, unknown>
         )
-      })
-      .then(() =>
-        this.options.onSuccess?.(
-          data,
-          this.state.variables!,
-          this.state.context!
-        )
+        const context = await this.options.onMutate?.(this.state.variables!)
+        if (context !== this.state.context) {
+          this.dispatch({
+            type: 'loading',
+            context,
+            variables: this.state.variables,
+          })
+        }
+      }
+      const data = await executeMutation()
+
+      // Notify cache callback
+      this.mutationCache.config.onSuccess?.(
+        data,
+        this.state.variables,
+        this.state.context,
+        this as Mutation<unknown, unknown, unknown, unknown>
       )
-      .then(() =>
-        this.options.onSettled?.(
-          data,
-          null,
-          this.state.variables!,
-          this.state.context
-        )
+
+      await this.options.onSuccess?.(
+        data,
+        this.state.variables!,
+        this.state.context!
       )
-      .then(() => {
-        this.dispatch({ type: 'success', data })
-        return data
-      })
-      .catch(error => {
+
+      await this.options.onSettled?.(
+        data,
+        null,
+        this.state.variables!,
+        this.state.context
+      )
+
+      this.dispatch({ type: 'success', data })
+      return data
+    } catch (error) {
+      try {
         // Notify cache callback
         this.mutationCache.config.onError?.(
           error,
@@ -258,27 +243,23 @@ export class Mutation<
           this.logger.error(error)
         }
 
-        return Promise.resolve()
-          .then(() =>
-            this.options.onError?.(
-              error,
-              this.state.variables!,
-              this.state.context
-            )
-          )
-          .then(() =>
-            this.options.onSettled?.(
-              undefined,
-              error,
-              this.state.variables!,
-              this.state.context
-            )
-          )
-          .then(() => {
-            this.dispatch({ type: 'error', error })
-            throw error
-          })
-      })
+        await this.options.onError?.(
+          error as TError,
+          this.state.variables!,
+          this.state.context
+        )
+
+        await this.options.onSettled?.(
+          undefined,
+          error as TError,
+          this.state.variables!,
+          this.state.context
+        )
+        throw error
+      } finally {
+        this.dispatch({ type: 'error', error: error as TError })
+      }
+    }
   }
 
   private dispatch(action: Action<TData, TError, TVariables, TContext>): void {

src/reactjs/tests/useMutation.test.tsx

@@ -855,4 +855,115 @@ describe('useMutation', () => {
       undefined
     )
   })
+
+  test('should go to error state if onSuccess callback errors', async () => {
+    const error = new Error('error from onSuccess')
+    const onError = jest.fn()
+
+    function Page() {
+      const mutation = useMutation(
+        async (_text: string) => {
+          await sleep(10)
+          return 'result'
+        },
+        {
+          onSuccess: () => Promise.reject(error),
+          onError,
+        }
+      )
+
+      return (
+        <div>
+          <button onClick={() => mutation.mutate('todo')}>mutate</button>
+          <div>status: {mutation.status}</div>
+        </div>
+      )
+    }
+
+    const rendered = renderWithClient(queryClient, <Page />)
+
+    await rendered.findByText('status: idle')
+
+    rendered.getByRole('button', { name: /mutate/i }).click()
+
+    await rendered.findByText('status: error')
+
+    expect(onError).toHaveBeenCalledWith(error, 'todo', undefined)
+  })
+
+  test('should go to error state if onError callback errors', async () => {
+    const error = new Error('error from onError')
+    const mutateFnError = new Error('mutateFnError')
+
+    function Page() {
+      const mutation = useMutation(
+        async (_text: string) => {
+          await sleep(10)
+          throw mutateFnError
+        },
+        {
+          onError: () => Promise.reject(error),
+        }
+      )
+
+      return (
+        <div>
+          <button onClick={() => mutation.mutate('todo')}>mutate</button>
+          <div>
+            error:{' '}
+            {mutation.error instanceof Error ? mutation.error.message : 'null'},
+            status: {mutation.status}
+          </div>
+        </div>
+      )
+    }
+
+    const rendered = renderWithClient(queryClient, <Page />)
+
+    await rendered.findByText('error: null, status: idle')
+
+    rendered.getByRole('button', { name: /mutate/i }).click()
+
+    await rendered.findByText('error: mutateFnError, status: error')
+  })
+
+  test('should go to error state if onSettled callback errors', async () => {
+    const error = new Error('error from onSettled')
+    const mutateFnError = new Error('mutateFnError')
+    const onError = jest.fn()
+
+    function Page() {
+      const mutation = useMutation(
+        async (_text: string) => {
+          await sleep(10)
+          throw mutateFnError
+        },
+        {
+          onSettled: () => Promise.reject(error),
+          onError,
+        }
+      )
+
+      return (
+        <div>
+          <button onClick={() => mutation.mutate('todo')}>mutate</button>
+          <div>
+            error:{' '}
+            {mutation.error instanceof Error ? mutation.error.message : 'null'},
+            status: {mutation.status}
+          </div>
+        </div>
+      )
+    }
+
+    const rendered = renderWithClient(queryClient, <Page />)
+
+    await rendered.findByText('error: null, status: idle')
+
+    rendered.getByRole('button', { name: /mutate/i }).click()
+
+    await rendered.findByText('error: mutateFnError, status: error')
+
+    expect(onError).toHaveBeenCalledWith(mutateFnError, 'todo', undefined)
+  })
 })