diff --git a/packages/store/src/atom.ts b/packages/store/src/atom.ts index 1a29e607..b5d41592 100644 --- a/packages/store/src/atom.ts +++ b/packages/store/src/atom.ts @@ -67,10 +67,10 @@ let queuedEffectsLength = 0 let activeSub: ReactiveNode | undefined let batchDepth = 0 -export function batch(fn: () => void) { +export function batch(fn: () => T): T { try { ++batchDepth - fn() + return fn() } finally { if (!--batchDepth) { flush() diff --git a/packages/store/tests/batch.test.ts b/packages/store/tests/batch.test.ts index 5550755a..4d9a51b7 100644 --- a/packages/store/tests/batch.test.ts +++ b/packages/store/tests/batch.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, test } from 'vitest' +import { describe, expect, test, vi } from 'vitest' import { batch, createStore } from '../src' describe('batch', () => { @@ -16,4 +16,26 @@ describe('batch', () => { expect(store.state.hello).toBe(101) expect(store.state.world).toBe(201) }) + + test('notifies subscribers once after entire batch completes', () => { + const store = createStore(0) + const listener = vi.fn() + store.subscribe(listener) + + batch(() => { + store.setState((prev) => prev + 1) + store.setState((prev) => prev + 1) + store.setState((prev) => prev + 1) + expect(listener).not.toHaveBeenCalled() + }) + + expect(store.state).toBe(3) + expect(listener).toHaveBeenCalledTimes(1) + expect(listener).toHaveBeenCalledWith(3) + }) + + test('batch returns the value produced by callback', () => { + const result = batch(() => 'batched-result') + expect(result).toBe('batched-result') + }) })