From 3e9d50307b4b63270dd12e87a29c6ac371be165e Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Thu, 16 Jan 2025 16:56:36 -0800 Subject: [PATCH] feat: observe api for observing stores with effects --- src/observe.ts | 30 +++++++ tests/observe.test.ts | 179 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 209 insertions(+) create mode 100644 src/observe.ts create mode 100644 tests/observe.test.ts diff --git a/src/observe.ts b/src/observe.ts new file mode 100644 index 0000000..d4c3f1e --- /dev/null +++ b/src/observe.ts @@ -0,0 +1,30 @@ +import { getDefaultStore } from 'jotai/vanilla' +import { Effect, atomEffect } from './atomEffect' + +type Store = ReturnType + +const storeEffects = new WeakMap void>>() + +type Unobserve = () => void + +export function observe(effect: Effect, store = getDefaultStore()): Unobserve { + let effectSubscriptions = storeEffects.get(store) + if (!effectSubscriptions) { + effectSubscriptions = new Map() + storeEffects.set(store, effectSubscriptions) + } + let unsubscribe = effectSubscriptions!.get(effect) + if (!unsubscribe) { + unsubscribe = store.sub(atomEffect(effect), () => void 0) + effectSubscriptions?.set(effect, unsubscribe) + } + return function unobserve() { + if (effectSubscriptions?.has(effect)) { + unsubscribe?.() + effectSubscriptions?.delete(effect) + if (effectSubscriptions.size === 0) { + storeEffects.delete(store) + } + } + } +} diff --git a/tests/observe.test.ts b/tests/observe.test.ts new file mode 100644 index 0000000..82dd19f --- /dev/null +++ b/tests/observe.test.ts @@ -0,0 +1,179 @@ +import { Getter, atom, createStore, getDefaultStore } from 'jotai/vanilla' +import { describe, expect, it } from 'vitest' +import { observe } from '../src/observe' +import { delay, increment } from './test-utils' + +describe('observe', () => { + it('should run effect on subscription and cleanup on unsubscribe', async () => { + const countAtom = atom(0) + let mounted = 0 + + const unsubscribe = observe((get, _set) => { + mounted++ + get(countAtom) + return () => { + mounted-- + } + }) + + await delay(0) + expect(mounted).toBe(1) + unsubscribe() + expect(mounted).toBe(0) + }) + + it('should reuse existing subscription for the same effect', async () => { + const countAtom = atom(0) + let runCount = 0 + + const effect = (get: any) => { + runCount++ + get(countAtom) + } + const store = getDefaultStore() + + const unsubscribe1 = observe(effect, store) + const unsubscribe2 = observe(effect, store) + + await delay(0) + expect(runCount).toBe(1) + + store.set(countAtom, increment) + await delay(0) + expect(runCount).toBe(2) + unsubscribe1() + unsubscribe2() + store.set(countAtom, increment) + await delay(0) + expect(runCount).toBe(2) + }) + + it('should unsubscribe the effect when any of effect subscriptions are unsubscribed', async () => { + const countAtom = atom(0) + let runCount = 0 + + const effect = (get: any) => { + runCount++ + get(countAtom) + } + const store = getDefaultStore() + + const unsubscribe1 = observe(effect, store) + const unsubscribe2 = observe(effect, store) + + await delay(0) + expect(runCount).toBe(1) + + store.set(countAtom, increment) + await delay(0) + expect(runCount).toBe(2) + unsubscribe1() + store.set(countAtom, increment) + await delay(0) + expect(runCount).toBe(2) + unsubscribe2() + expect(runCount).toBe(2) + store.set(countAtom, increment) + await delay(0) + expect(runCount).toBe(2) + }) + + it('should work with custom store', async () => { + const store = createStore() + const countAtom = atom(0) + let runCount = 0 + + const unsubscribe = observe((get) => { + runCount++ + get(countAtom) + }, store) + + await delay(0) + expect(runCount).toBe(1) + + store.set(countAtom, increment) + await delay(0) + expect(runCount).toBe(2) + + unsubscribe() + }) + + it('should handle multiple stores independently', async () => { + const store1 = createStore() + const store2 = createStore() + const countAtom = atom(0) + const runCounts = [0, 0] as [number, number] + + const storeIdAtom = atom(NaN as 0 | 1) + store1.set(storeIdAtom, 0) + store2.set(storeIdAtom, 1) + + function effect(get: Getter) { + ++runCounts[get(storeIdAtom)] + get(countAtom) + } + + const unsubscribe1 = observe(effect, store1) + const unsubscribe2 = observe(effect, store2) + + await delay(0) + expect(runCounts[0]).toBe(1) + expect(runCounts[1]).toBe(1) + + store1.set(countAtom, increment) + await delay(0) + expect(runCounts[0]).toBe(2) + expect(runCounts[1]).toBe(1) + store2.set(countAtom, increment) + await delay(0) + expect(runCounts[0]).toBe(2) + expect(runCounts[1]).toBe(2) + + unsubscribe1() + store2.set(countAtom, increment) + await delay(0) + expect(runCounts[0]).toBe(2) + expect(runCounts[1]).toBe(3) + store2.set(countAtom, increment) + await delay(0) + expect(runCounts[0]).toBe(2) + expect(runCounts[1]).toBe(4) + + unsubscribe2() + }) + + it('should handle multiple effects independently', async () => { + const store = createStore() + const countAtom = atom(0) + const runCounts = [0, 0] as [number, number] + + function effect1(get: Getter) { + ++runCounts[0] + get(countAtom) + } + function effect2(get: Getter) { + ++runCounts[1] + get(countAtom) + } + + const unsubscribe1 = observe(effect1, store) + const unsubscribe2 = observe(effect2, store) + + await delay(0) + expect(runCounts[0]).toBe(1) + expect(runCounts[1]).toBe(1) + + store.set(countAtom, increment) + await delay(0) + expect(runCounts[0]).toBe(2) + expect(runCounts[1]).toBe(2) + + unsubscribe1() + store.set(countAtom, increment) + await delay(0) + expect(runCounts[0]).toBe(2) + expect(runCounts[1]).toBe(3) + + unsubscribe2() + }) +})