all working

__tests__/atomEffect.strict.test.ts

@@ -1,12 +1,12 @@
 import { StrictMode, useEffect } from 'react'
 import { act, renderHook, waitFor } from '@testing-library/react'
 import { useAtom, useAtomValue, useSetAtom } from 'jotai/react'
-import { atom } from 'jotai/vanilla'
+import { atom, getDefaultStore } from 'jotai/vanilla'
 import { atomEffect } from '../src/atomEffect'
 
 const wrapper = StrictMode
 
-it.only('should run the effect on mount and cleanup on unmount once', async () => {
+it('should run the effect on mount and cleanup on unmount once', async () => {
   expect.assertions(5)
   const effect = { mount: 0, unmount: 0 }
 
@@ -108,32 +108,34 @@ it('should run the effect on mount and cleanup on unmount and whenever countAtom
 })
 
 it('should not cause infinite loops when effect updates the watched atom', async () => {
-  expect.assertions(2)
+  expect.assertions(1)
   const watchedAtom = atom(0)
   let runCount = 0
   const effectAtom = atomEffect((get, set) => {
     get(watchedAtom)
     runCount++
-    set(watchedAtom, get(watchedAtom) + 1)
+    console.log('incrementing watchedAtom...')
+    set(watchedAtom, increment)
+    return () => {
+      set(watchedAtom, (c) => c - 1)
+    }
   })
+  const store = getDefaultStore()
   function useTest() {
-    useAtom(effectAtom)
-    const setCount = useSetAtom(watchedAtom)
-    return () => act(async () => setCount(increment))
+    useAtom(effectAtom, { store })
   }
-  const { result, rerender } = renderHook(useTest, { wrapper: StrictMode })
+  const { rerender } = renderHook(useTest, { wrapper: StrictMode })
 
   // initial render should run the effect once
   await waitFor(() => assert(runCount === 1))
-
-  rerender()
-
   // rerender should not run the effect again
-  expect(runCount).toBe(1)
+  rerender()
+  await delay(0)
 
-  // changing the value should run the effect again one time
-  await result.current()
-  expect(runCount).toBe(2)
+  expect({ runCount, watched: store.get(watchedAtom) }).toEqual({
+    runCount: 1,
+    watched: 1,
+  })
 })
 
 it('should not cause infinite loops when effect updates the watched atom asynchronous', async () => {
@@ -163,9 +165,34 @@ it('should not cause infinite loops when effect updates the watched atom asynchr
   expect(runCount).toBe(2)
 })
 
-////////////////////////////////////////////////////////////////////////////////////////////
-
-////////////////////////////////////////////////////////////////////////////////////////////
+it.only('should allow synchronous infinite loops with opt-in for first run', async () => {
+  expect.assertions(1)
+  let runCount = 0
+  const watchedAtom = atom(0)
+  const effectAtom = atomEffect((get, { recurse }) => {
+    const value = get(watchedAtom)
+    runCount++
+    if (value < 5) {
+      console.log('setting watchedAtom...')
+      console.log('in effect: recurse', 'watchedAtom:', get(watchedAtom))
+      recurse(watchedAtom, increment)
+    }
+  })
+  const store = getDefaultStore()
+  function useTest() {
+    useAtom(effectAtom, { store })
+    const setCount = useSetAtom(watchedAtom, { store })
+    return () => act(async () => setCount(increment))
+  }
+  const { result } = renderHook(useTest, { wrapper })
+  await delay(0)
+  await act(async () => result.current())
+  await delay(0)
+  expect({ runCount, watched: store.get(watchedAtom) }).toEqual({
+    runCount: 7, // extra run for strict mode render
+    watched: 6,
+  })
+})
 it('should conditionally run the effect and cleanup when effectAtom is unmounted', async () => {
   expect.assertions(6)
 

__tests__/atomEffect.test.ts

@@ -184,25 +184,31 @@ it('should not cause infinite loops when effect updates the watched atom asynchr
   await delay(0)
   expect(runCount).toBe(2)
 })
-//////////////////////////////////////////////////////////////////
+
 it('should allow synchronous infinite loops with opt-in for first run', async () => {
-  expect.assertions(2)
+  expect.assertions(1)
   let runCount = 0
   const watchedAtom = atom(0)
+  let done = false
   const effectAtom = atomEffect((get, { recurse }) => {
     const value = get(watchedAtom)
-    runCount++
-    if (value < 5) {
-      console.log('setting watchedAtom...')
-      console.log('in effect: recurse', 'watchedAtom:', get(watchedAtom))
-      recurse(watchedAtom, increment)
+    const thisRunCount = runCount++
+    console.log('in effect', { thisRunCount, value })
+    if (value >= 3) {
+      console.red('GTE 3, EXITING --------')
+      done = true
+      return
     }
+    console.log('setting watchedAtom...', { value })
+    recurse(watchedAtom, increment)
   })
   const store = getDefaultStore()
   store.sub(effectAtom, () => void 0)
-  await waitFor(() => assert(store.get(watchedAtom) === 5))
-  expect(store.get(watchedAtom)).toBe(5)
-  expect(runCount).toBe(6)
+  await waitFor(() => assert(done))
+  expect({ runCount, watched: store.get(watchedAtom) }).toEqual({
+    runCount: 4,
+    watched: 3,
+  })
 })
 
 it('should allow synchronous infinite loops with opt-in', async () => {
@@ -415,24 +421,75 @@ it('should work with both recurse and set', async () => {
     const value = get(watchedAtom)
     get(countAtom)
     runCount++
-    if (value === 0 || value % 5) {
+    if (value === 0 || value % 3) {
       set.recurse(watchedAtom, increment)
       set(countAtom, increment)
       console.log(get(watchedAtom), get(countAtom))
-    } else {
-      console.log('----- sole increment -----')
-      set(watchedAtom, increment)
+      return
     }
-    console.log('setting watchedAtom complete')
+    console.log('----- sole increment -----')
+    set(watchedAtom, increment)
   })
   const store = getDefaultStore()
   store.sub(effectAtom, () => void 0)
-  await waitFor(() => assert(store.get(countAtom) === 5))
-  expect(store.get(countAtom)).toBe(5)
-  expect(store.get(watchedAtom)).toBe(6)
-  expect(runCount).toBe(6)
+  await waitFor(() => assert(store.get(countAtom) === 3))
+  expect(store.get(countAtom)).toBe(3)
+  expect(store.get(watchedAtom)).toBe(4)
+  expect(runCount).toBe(4)
+})
+
+it('should disallow synchronous infinite loops in cleanup', async () => {
+  expect.assertions(2)
+  let runCount = 0
+  const watchedAtom = atom(0)
+  const anotherAtom = atom(0)
+  const effectAtom = atomEffect((get, { recurse }) => {
+    const value = get(watchedAtom)
+    get(anotherAtom)
+    const thisRunCount = runCount++
+    console.log('in effect', { thisRunCount, value })
+    return () => {
+      console.log('in cleanup', { thisRunCount, value })
+      recurse(watchedAtom, increment)
+    }
+  })
+  const store = getDefaultStore()
+  console.yellow('mounting effect')
+  store.sub(effectAtom, () => void 0)
+  await delay(0)
+  store.set(anotherAtom, increment)
+  await delay(0)
+  console.yellow('increment watched')
+  console.log('final watchedAtom:', store.get(watchedAtom))
+  expect(store.get(watchedAtom)).toBe(0)
+  expect(runCount).toBe(2)
+})
+
+// FIXME: is there a way to disallow asynchronous infinite loops in cleanup?
+
+it('should return value from recurse', async () => {
+  expect.assertions(1)
+  const countAtom = atom(0)
+  const incrementCountAtom = atom(null, (get, set) => {
+    set(countAtom, increment)
+    return get(countAtom)
+  })
+  const results = [] as number[]
+  let done = false
+  const effectAtom = atomEffect((get, { recurse }) => {
+    const value = get(countAtom)
+    if (value < 5) {
+      const result = recurse(incrementCountAtom)
+      results.unshift(result)
+      done = true
+      return
+    }
+  })
+  const store = getDefaultStore()
+  store.sub(effectAtom, () => void 0)
+  await waitFor(() => assert(done))
+  expect(results).toEqual([1, 2, 3, 4, 5])
 })
-//////////////////////////////////////////////////////////////////
 
 it('should conditionally run the effect and cleanup when effectAtom is unmounted', async () => {
   expect.assertions(6)

src/atomEffect.ts

@@ -14,20 +14,19 @@ export function atomEffect(
     inProgress: 0,
     promise: undefined as Promise<void> | undefined,
     cleanup: undefined as CleanupFn | void,
+    fromCleanup: false,
     recursing: false,
     refresh: () => {},
     refreshing: false,
     set: (() => {}) as Setter,
   }))
   if (process.env.NODE_ENV !== 'production') {
     refAtom.debugPrivate = true
-    refAtom.debugLabel = 'refAtom' // remove
   }
 
   const refreshAtom = atom(0)
   if (process.env.NODE_ENV !== 'production') {
     refreshAtom.debugPrivate = true
-    refreshAtom.debugLabel = 'refreshAtom' // remove
   }
 
   const initAtom = atom(null, (get, set, mounted: boolean) => {
@@ -43,7 +42,7 @@ export function atomEffect(
       }
     }
     if (mounted) {
-      ref.refresh()
+      set(refreshAtom, (c) => c + 1)
     } else {
       ref.cleanup?.()
       ref.cleanup = undefined
@@ -59,21 +58,22 @@ export function atomEffect(
   }
   if (process.env.NODE_ENV !== 'production') {
     initAtom.debugPrivate = true
-    initAtom.debugLabel = 'initAtom' // remove
   }
 
   const effectAtom = atom((get) => {
     const refreshValue = get(refreshAtom)
     const ref = get(refAtom)
     if (!ref.mounted) {
+      console.red('[not mounted, returning...]')
       return ref.promise
     }
-    const { inProgress, recursing, refreshing } = ref
+    const { fromCleanup, inProgress, recursing, refreshing } = ref
     console.lightGreen('[atomEffect]', {
       inProgress,
       recursing,
       refreshing,
       refreshValue,
+      fromCleanup,
     })
     if (ref.inProgress && !ref.refreshing) {
       console.red('[in progress, returning...]')
@@ -83,7 +83,7 @@ export function atomEffect(
       console.red('[in recursing, returning...]')
       return ref.promise
     }
-    console.lightBlue('[refreshing running...]')
+    console.lightBlue('[running effect...]')
     ref.refreshing = false
     const currDeps = new Map<Atom<unknown>, unknown>()
     const getter: Getter = (a) => {
@@ -100,6 +100,10 @@ export function atomEffect(
       }
     }
     setter.recurse = (anAtom, ...args) => {
+      if (ref.fromCleanup) {
+        console.warn('cannot recurse from cleanup')
+        return undefined as any
+      }
       ref.recursing = true
       try {
         console.yellow('[calling recurse...]')
@@ -127,9 +131,19 @@ export function atomEffect(
     const effector = async () => {
       try {
         if (!ref.mounted) {
+          console.red('[in effector: not mounted, returning...]')
           return
         }
-        ref.cleanup?.()
+        ref.fromCleanup = true
+        try {
+          ref.cleanup
+            ? console.orange('[running cleanup...]')
+            : console.orange('[no cleanup]')
+          ref.cleanup?.()
+        } finally {
+          ref.fromCleanup = false
+        }
+        console.orange('[running effectFn...]')
         ref.cleanup = effectFn(getter, setter)
       } finally {
         --ref.inProgress