From 924f7ff68696b69c716f733fdede8ae479053791 Mon Sep 17 00:00:00 2001
From: grzdev <grzdev@users.noreply.github.com>
Date: Sun, 11 Oct 2026 19:53:28 +0100
Subject: [PATCH] feat(store): return callback result from batch and add
 subscriber test coverage

---
 packages/store/src/atom.ts         |  4 ++--
 packages/store/tests/batch.test.ts | 24 +++++++++++++++++++++++-
 2 files changed, 25 insertions(+), 3 deletions(-)

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<T>(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')
+  })
 })