feat: add dev_subscribe_store (pmndrs#1790)

* feat: add dev_subscribe_store * attach notifiers on public apis * revert changes * notify on set * notify mount and unmount * fix: update unmount + state + backwards compatibility for dev_subscribe_store * fix store listeners, and deperecation message for state listener * recover original dev_subscribe_state * test: add tests for dev_* methods * test: add [DEV-ONLY] tag to tests Co-authored-by: Daishi Kato <[email protected]> * chore: adjust sed script for dev/prd only tests * test: fix test description * adding tests * test: add failing test to unmount tree dependencies on unsub * fix test styles * do not fire before flush pending --------- Co-authored-by: daishi <[email protected]> Co-authored-by: Daishi Kato <[email protected]>
edisonLzy · Mar 2, 2023 · e80be16 · e80be16
1 parent cb54160
commit e80be16
Show file tree

Hide file tree

Showing 3 changed files with 172 additions and 6 deletions.
diff --git a/.github/workflows/test-multiple-builds.yml b/.github/workflows/test-multiple-builds.yml
@@ -30,13 +30,13 @@ jobs:
       - name: Patch for DEV-ONLY
         if: ${{ matrix.env == 'development' }}
         run: |
-          sed -i~ "s/it[.a-zA-Z]*('\[DEV-ONLY\]/it('/" tests/*/*.tsx tests/*/*/*.tsx
-          sed -i~ "s/it[.a-zA-Z]*('\[PRD-ONLY\]/it.skip('/" tests/*/*.tsx tests/*/*/*.tsx
+          sed -i~ "s/\(it\|describe\)[.a-zA-Z]*('\[DEV-ONLY\]/\1('/" tests/*/*.tsx tests/*/*/*.tsx
+          sed -i~ "s/\(it\|describe\)[.a-zA-Z]*('\[PRD-ONLY\]/\1.skip('/" tests/*/*.tsx tests/*/*/*.tsx
       - name: Patch for PRD-ONLY
         if: ${{ matrix.env == 'production' }}
         run: |
-          sed -i~ "s/it[.a-zA-Z]*('\[PRD-ONLY\]/it('/" tests/*/*.tsx tests/*/*/*.tsx
-          sed -i~ "s/it[.a-zA-Z]*('\[DEV-ONLY\]/it.skip('/" tests/*/*.tsx tests/*/*/*.tsx
+          sed -i~ "s/\(it\|describe\)[.a-zA-Z]*('\PRD-ONLY\]/\1('/" tests/*/*.tsx tests/*/*/*.tsx
+          sed -i~ "s/\(it\|describe\)[.a-zA-Z]*('\[DEV-ONLY\]/\1.skip('/" tests/*/*.tsx tests/*/*/*.tsx
       - name: Patch for CJS
         if: ${{ matrix.build == 'cjs' }}
         run: |

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
@@ -109,6 +109,7 @@ type Mounted = {
 
 // for debugging purpose only
 type StateListener = () => void
+type StoreListener = (type: 'state' | 'sub' | 'unsub') => void
 type MountedAtoms = Set<AnyAtom>
 
 /**
@@ -135,9 +136,11 @@ export const createStore = () => {
     AtomState /* prevAtomState */ | undefined
   >()
   let stateListeners: Set<StateListener>
+  let storeListeners: Set<StoreListener>
   let mountedAtoms: MountedAtoms
   if (import.meta.env?.MODE !== 'production') {
     stateListeners = new Set()
+    storeListeners = new Set()
     mountedAtoms = new Set()
   }
 
@@ -582,6 +585,7 @@ export const createStore = () => {
     }
     if (import.meta.env?.MODE !== 'production') {
       stateListeners.forEach((l) => l())
+      storeListeners.forEach((l) => l('state'))
     }
   }
 
@@ -590,9 +594,16 @@ export const createStore = () => {
     flushPending()
     const listeners = mounted.l
     listeners.add(listener)
+    if (import.meta.env?.MODE !== 'production') {
+      storeListeners.forEach((l) => l('sub'))
+    }
     return () => {
       listeners.delete(listener)
       delAtom(atom)
+      if (import.meta.env?.MODE !== 'production') {
+        // devtools uses this to detect if it _can_ unmount or not
+        storeListeners.forEach((l) => l('unsub'))
+      }
     }
   }
 
@@ -601,13 +612,22 @@ export const createStore = () => {
       get: readAtom,
       set: writeAtom,
       sub: subscribeAtom,
-      // store dev methods (these are tentative and subject to change)
+      // store dev methods (these are tentative and subject to change without notice)
       dev_subscribe_state: (l: StateListener) => {
+        console.warn(
+          '[DEPRECATED] dev_subscribe_state is deprecated and will be removed in the next minor version. use dev_subscribe_store instead.'
+        )
         stateListeners.add(l)
         return () => {
           stateListeners.delete(l)
         }
       },
+      dev_subscribe_store: (l: StoreListener) => {
+        storeListeners.add(l)
+        return () => {
+          storeListeners.delete(l)
+        }
+      },
       dev_get_mounted_atoms: () => mountedAtoms.values(),
       dev_get_atom_state: (a: AnyAtom) => atomStateMap.get(a),
       dev_get_mounted: (a: AnyAtom) => mountedMap.get(a),

diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx
@@ -1,4 +1,4 @@
-import { expect, it, jest } from '@jest/globals'
+import { describe, expect, it, jest } from '@jest/globals'
 import { atom, createStore } from 'jotai/vanilla'
 
 it('should not fire on subscribe', async () => {
@@ -32,3 +32,149 @@ it('should not fire subscription if derived atom value is the same', async () =>
   store.set(countAtom, 1)
   expect(callback).toBeCalledTimes(calledTimes)
 })
+
+describe('[DEV-ONLY] dev-only methods', () => {
+  it('should return the values of all mounted atoms', () => {
+    const store = createStore()
+    const countAtom = atom(0)
+    countAtom.debugLabel = 'countAtom'
+    const derivedAtom = atom((get) => get(countAtom) * 0)
+    const unsub = store.sub(derivedAtom, jest.fn())
+    store.set(countAtom, 1)
+
+    const result = store.dev_get_mounted_atoms?.() || []
+    expect(Array.from(result)).toStrictEqual([
+      { toString: expect.any(Function), read: expect.any(Function) },
+      {
+        toString: expect.any(Function),
+        init: 0,
+        read: expect.any(Function),
+        write: expect.any(Function),
+        debugLabel: 'countAtom',
+      },
+    ])
+    unsub()
+  })
+
+  it('should get atom state of a given atom', () => {
+    const store = createStore()
+    const countAtom = atom(0)
+    countAtom.debugLabel = 'countAtom'
+    const unsub = store.sub(countAtom, jest.fn())
+    store.set(countAtom, 1)
+    const result = store.dev_get_atom_state?.(countAtom)
+    expect(result).toHaveProperty('v', 1)
+    unsub()
+  })
+
+  it('should get mounted atom from mounted map', () => {
+    const store = createStore()
+    const countAtom = atom(0)
+    countAtom.debugLabel = 'countAtom'
+    const cb = jest.fn()
+    const unsub = store.sub(countAtom, cb)
+    store.set(countAtom, 1)
+    const result = store.dev_get_mounted?.(countAtom)
+    expect(result).toStrictEqual({ l: new Set([cb]), t: new Set([countAtom]) })
+    unsub()
+  })
+
+  it('should restore atoms and its dependencies correctly', () => {
+    const store = createStore()
+    const countAtom = atom(0)
+    countAtom.debugLabel = 'countAtom'
+    const derivedAtom = atom((get) => get(countAtom) * 2)
+    store.set(countAtom, 1)
+    store.dev_restore_atoms?.([[countAtom, 2]])
+    expect(store.get(countAtom)).toBe(2)
+    expect(store.get?.(derivedAtom)).toBe(4)
+  })
+
+  describe('dev_subscribe_state', () => {
+    it('should call the callback when state change is flushed out', () => {
+      const store = createStore()
+      const callback = jest.fn()
+      const unsub = store.dev_subscribe_state?.(callback)
+      const countAtom = atom(0)
+      const unsubAtom = store.sub(countAtom, jest.fn())
+      expect(callback).toHaveBeenCalledTimes(1)
+      unsub?.()
+      unsubAtom?.()
+    })
+  })
+
+  describe('dev_subscribe_store', () => {
+    it('should call the callback when state changes', () => {
+      const store = createStore()
+      const callback = jest.fn()
+      const unsub = store.dev_subscribe_store?.(callback)
+      const countAtom = atom(0)
+      const unsubAtom = store.sub(countAtom, jest.fn())
+      store.set(countAtom, 1)
+      expect(callback).toHaveBeenNthCalledWith(1, 'state')
+      expect(callback).toHaveBeenNthCalledWith(2, 'sub')
+      expect(callback).toHaveBeenNthCalledWith(3, 'state')
+      expect(callback).toHaveBeenCalledTimes(3)
+      unsub?.()
+      unsubAtom?.()
+    })
+
+    it('should call unsub only when atom is unsubscribed', () => {
+      const store = createStore()
+      const callback = jest.fn()
+      const unsub = store.dev_subscribe_store?.(callback)
+      const countAtom = atom(0)
+      const unsubAtom = store.sub(countAtom, jest.fn())
+      const unsubAtomSecond = store.sub(countAtom, jest.fn())
+      unsubAtom?.()
+      expect(callback).toHaveBeenNthCalledWith(1, 'state')
+      expect(callback).toHaveBeenNthCalledWith(2, 'sub')
+      expect(callback).toHaveBeenNthCalledWith(3, 'state')
+      expect(callback).toHaveBeenNthCalledWith(4, 'sub')
+      expect(callback).toHaveBeenNthCalledWith(5, 'unsub')
+      expect(callback).toHaveBeenCalledTimes(5)
+      unsub?.()
+      unsubAtomSecond?.()
+    })
+  })
+})
+
+it('should unmount with store.get', async () => {
+  const store = createStore()
+  const countAtom = atom(0)
+  const callback = jest.fn()
+  const unsub = store.sub(countAtom, callback)
+  store.get(countAtom)
+  unsub()
+  const result = Array.from(store.dev_get_mounted_atoms?.() ?? [])
+  expect(result).toEqual([])
+})
+
+it('should unmount dependencies with store.get', async () => {
+  const store = createStore()
+  const countAtom = atom(0)
+  const derivedAtom = atom((get) => get(countAtom) * 2)
+  const callback = jest.fn()
+  const unsub = store.sub(derivedAtom, callback)
+  store.get(derivedAtom)
+  unsub()
+  const result = Array.from(store.dev_get_mounted_atoms?.() ?? [])
+  expect(result).toEqual([])
+})
+
+it('should unmount tree dependencies with store.get', async () => {
+  const store = createStore()
+  const countAtom = atom(0)
+  const derivedAtom = atom((get) => get(countAtom) * 2)
+  const anotherDerivedAtom = atom((get) => get(countAtom) * 3)
+  const callback = jest.fn()
+  const unsubStore = store.dev_subscribe_store?.(() => {
+    // Comment this line to make the test pass
+    store.get(derivedAtom)
+  })
+  const unsub = store.sub(anotherDerivedAtom, callback)
+  unsub()
+  unsubStore?.()
+  const result = Array.from(store.dev_get_mounted_atoms?.() ?? [])
+  expect(result).toEqual([])
+})