allow atom with default update function to return reset and add tests (#3137)

* allow atom with default update function to return reset and add tests

* dont export type, and clarify contract of dynamic default with a few more tests

---------

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

Author: scamden

src/vanilla/utils/atomWithDefault.ts

@@ -1,5 +1,5 @@
 import { atom } from '../../vanilla.ts'
-import type { SetStateAction, WritableAtom } from '../../vanilla.ts'
+import type { WritableAtom } from '../../vanilla.ts'
 import { RESET } from './constants.ts'
 
 type Read<Value, Args extends unknown[], Result> = WritableAtom<
@@ -8,38 +8,37 @@ type Read<Value, Args extends unknown[], Result> = WritableAtom<
   Result
 >['read']
 
+type DefaultSetStateAction<Value> =
+  | Value
+  | typeof RESET
+  | ((prev: Value) => Value | typeof RESET)
+
 export function atomWithDefault<Value>(
-  getDefault: Read<Value, [SetStateAction<Value> | typeof RESET], void>,
-): WritableAtom<Value, [SetStateAction<Value> | typeof RESET], void> {
+  getDefault: Read<Value, [DefaultSetStateAction<Value>], void>,
+): WritableAtom<Value, [DefaultSetStateAction<Value>], void> {
   const EMPTY = Symbol()
   const overwrittenAtom = atom<Value | typeof EMPTY>(EMPTY)
 
   if (import.meta.env?.MODE !== 'production') {
     overwrittenAtom.debugPrivate = true
   }
 
-  const anAtom: WritableAtom<
-    Value,
-    [SetStateAction<Value> | typeof RESET],
-    void
-  > = atom(
-    (get, options) => {
-      const overwritten = get(overwrittenAtom)
-      if (overwritten !== EMPTY) {
-        return overwritten
-      }
-      return getDefault(get, options)
-    },
-    (get, set, update) => {
-      if (update === RESET) {
-        set(overwrittenAtom, EMPTY)
-      } else if (typeof update === 'function') {
-        const prevValue = get(anAtom)
-        set(overwrittenAtom, (update as (prev: Value) => Value)(prevValue))
-      } else {
-        set(overwrittenAtom, update)
-      }
-    },
-  )
+  const anAtom: WritableAtom<Value, [DefaultSetStateAction<Value>], void> =
+    atom(
+      (get, options) => {
+        const overwritten = get(overwrittenAtom)
+        if (overwritten !== EMPTY) {
+          return overwritten
+        }
+        return getDefault(get, options)
+      },
+      (get, set, update) => {
+        const newValue =
+          typeof update === 'function'
+            ? (update as (prev: Value) => Value)(get(anAtom))
+            : update
+        set(overwrittenAtom, newValue === RESET ? EMPTY : newValue)
+      },
+    )
   return anAtom
 }

tests/vanilla/utils/atomWithDefault.test.ts

@@ -0,0 +1,75 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { atom, createStore } from 'jotai/vanilla'
+import { RESET, atomWithDefault } from 'jotai/vanilla/utils'
+
+describe('atomWithDefault', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('should reset to initial value using RESET', () => {
+    const initialValue = 10
+    const testAtom = atomWithDefault(() => initialValue)
+    const store = createStore()
+    store.set(testAtom, 123)
+    store.set(testAtom, RESET)
+    expect(store.get(testAtom)).toBe(initialValue)
+  })
+
+  it('should reset to initial value derived from another atom', () => {
+    const initialValueAtom = atom(10)
+    const testAtom = atomWithDefault((get) => get(initialValueAtom))
+    const store = createStore()
+    expect(store.get(testAtom)).toBe(10)
+    store.set(testAtom, 123)
+    store.set(initialValueAtom, 20)
+    store.set(testAtom, RESET)
+    expect(store.get(testAtom)).toBe(20)
+  })
+
+  it(`should reflect changes to the initial value atom when main atom hasn't been manually changed`, () => {
+    const initialValueAtom = atom(10)
+    const testAtom = atomWithDefault((get) => get(initialValueAtom))
+    const store = createStore()
+    store.set(initialValueAtom, 20)
+    expect(store.get(testAtom)).toBe(20)
+  })
+
+  it(`should reflect changes to the initial value atom when main atom has been manually changed but then RESET`, () => {
+    const initialValueAtom = atom(10)
+    const testAtom = atomWithDefault((get) => get(initialValueAtom))
+    const store = createStore()
+    store.set(testAtom, 123)
+    // if this RESET were storing 10 rather than EMPTY the next set wouldn't have an effect
+    store.set(testAtom, RESET)
+    store.set(initialValueAtom, 20)
+    expect(store.get(testAtom)).toBe(20)
+  })
+
+  it('should update atom with a new value', () => {
+    const initialValue = 10
+    const testAtom = atomWithDefault(() => initialValue)
+    const store = createStore()
+    store.set(testAtom, 123)
+    store.set(testAtom, 30)
+    expect(store.get(testAtom)).toBe(30)
+  })
+
+  it('should update atom using a function', () => {
+    const initialValue = 10
+    const testAtom = atomWithDefault(() => initialValue)
+    const store = createStore()
+    store.set(testAtom, 123)
+    store.set(testAtom, (prev: number) => prev + 10)
+    expect(store.get(testAtom)).toBe(133)
+  })
+
+  it('should reset with a function', () => {
+    const initialValue = 10
+    const testAtom = atomWithDefault(() => initialValue)
+    const store = createStore()
+    store.set(testAtom, 123)
+    store.set(testAtom, () => RESET)
+    expect(store.get(testAtom)).toBe(initialValue)
+  })
+})