fix: loading mutation does not call refresh on unmount (#38)

* failing test

* fix

* fix

* lint

* remove extra line

* as unknown as never

* Update src/atomsWithMutation.ts

---------

Co-authored-by: Daishi Kato <[email protected]>

__tests__/atomWithMutation.tsx

@@ -0,0 +1,58 @@
+import React, { useState } from 'react'
+import { fireEvent, render } from '@testing-library/react'
+import { useAtom } from 'jotai/react'
+import { atomsWithMutation } from '../src/index'
+
+it('atomsWithMutation should be refreshed on unmount (#2060)', async () => {
+  let resolve: (() => void) | undefined
+  const [, testAtom] = atomsWithMutation<number, number, number, number>(
+    () => ({
+      mutationKey: ['test-atom'],
+      mutationFn: async (a) => {
+        await new Promise<void>((r) => {
+          resolve = r
+        })
+        return a
+      },
+    })
+  )
+
+  function App() {
+    const [mount, setMount] = useState<boolean>(true)
+    return (
+      <div>
+        <button onClick={() => setMount(false)}>unmount</button>
+        <button onClick={() => setMount(true)}>mount</button>
+        {mount && <TestView />}
+      </div>
+    )
+  }
+
+  function TestView() {
+    const [state, mutate] = useAtom(testAtom)
+    return (
+      <div>
+        <p>status: {state.status}</p>
+        <button disabled={state.isLoading} onClick={() => mutate([1])}>
+          mutate
+        </button>
+      </div>
+    )
+  }
+
+  const { findByText, getByText } = render(<App />)
+
+  await findByText('status: idle')
+
+  fireEvent.click(getByText('mutate'))
+  await findByText('status: loading')
+  resolve?.()
+  await findByText('status: success')
+
+  fireEvent.click(getByText('mutate'))
+  await findByText('status: loading')
+  fireEvent.click(getByText('unmount'))
+  fireEvent.click(getByText('mount'))
+  await findByText('status: idle')
+  resolve?.()
+})

src/common.ts

@@ -102,17 +102,22 @@ export const createAtoms = <
       const resultAtom = get(baseStatusAtom)
       return get(resultAtom)
     },
-    (get, set, action: Action) => {
+    (get, set, action: Action | 'refresh') => {
       const observer = get(observerAtom)
       const refresh = () => {
         const queryClient = getQueryClient(get)
         const observerCache = get(observerCacheAtom)
         observerCache.delete(queryClient)
         set(refreshAtom, (c) => c + 1)
       }
+      if (action === 'refresh') {
+        refresh()
+        return undefined as unknown as never
+      }
       return handleAction(action, observer, refresh)
     }
   )
+  statusAtom.onMount = (setAtom) => () => setAtom('refresh')
 
   const baseDataAtom = atom((get) => {
     getOptions(get) // re-create observable when options change