refactor initAtom.onMount and unskip tests

__tests__/atomEffect.test.tsx

@@ -820,8 +820,7 @@ it('should abort the previous promise', async () => {
   expect(completedRuns).toEqual([0, 2])
 })
 
-// TODO: enable this test when https://github.com/pmndrs/jotai/pull/2347 releases
-it.skip('should not infinite loop with nested atomEffects', async () => {
+it('should not infinite loop with nested atomEffects', async () => {
   const metrics = {
     mounted: 0,
     runCount1: 0,
@@ -913,9 +912,10 @@ it('should trigger the error boundary when an error is thrown', async () => {
   }
   render(<TestComponent />, { wrapper })
   await waitFor(() => assert(didThrow))
+  expect(didThrow).toBe(true)
 })
 
-it.only('should trigger an error boundary when an error is thrown in a cleanup', async () => {
+it('should trigger an error boundary when an error is thrown in a cleanup', async () => {
   expect.assertions(1)
 
   const refreshAtom = atom(0)

src/atomEffect.ts

@@ -24,33 +24,28 @@ export function atomEffect(
 
   const refreshAtom = atom(0)
 
-  const initAtom = atom(null, (get, set, mounted: boolean) => {
+  const initAtom = atom(null, (get, set) => {
     const ref = get(refAtom)
-    ref.mounted = mounted
-    if (mounted) {
-      ref.get = get
-      ref.set = set
-      ref.refresh = () => {
-        try {
-          ref.refreshing = true
-          set(refreshAtom, (c) => c + 1)
-        } finally {
-          ref.refreshing = false
-        }
+    ref.mounted = true
+    ref.get = get
+    ref.set = set
+    ref.refresh = () => {
+      try {
+        ref.refreshing = true
+        set(refreshAtom, (c) => c + 1)
+      } finally {
+        ref.refreshing = false
       }
-      set(refreshAtom, (c) => c + 1)
-    } else {
+    }
+    set(refreshAtom, (c) => c + 1)
+    return () => {
+      ref.mounted = false
       throwIfPendingError(ref)
       ref.cleanup?.()
       ref.cleanup = undefined
     }
   })
-  initAtom.onMount = (init) => {
-    init(true)
-    return () => {
-      init(false)
-    }
-  }
+  initAtom.onMount = (mount) => mount()
   const effectAtom = atom((get) => {
     get(refreshAtom)
     const ref = get(refAtom)