From 788a01f33a473dcb0a0847410fa62463a9b79a04 Mon Sep 17 00:00:00 2001
From: mesqueeb <mesqueeb@users.noreply.github.com>
Date: Tue, 4 Jun 2024 13:28:55 +0900
Subject: [PATCH] fix: allow `storeSplit` to be used in hooks

---
 .../src/moduleActions/handleWritePerStore.ts  |  4 +-
 .../test/internal/storeSplit.test.ts          | 40 +++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/packages/core/src/moduleActions/handleWritePerStore.ts b/packages/core/src/moduleActions/handleWritePerStore.ts
index 07cc202c..3993c965 100644
--- a/packages/core/src/moduleActions/handleWritePerStore.ts
+++ b/packages/core/src/moduleActions/handleWritePerStore.ts
@@ -122,8 +122,8 @@ export function handleWritePerStore(
         }>((dic, storeName) => ({ ...dic, [storeName]: unwrapStoreSplits(payload, storeName) }), {})
 
         for (const modifyFn of modifyPayloadFnsMap[actionName]) {
-          storePayloadDic = mapObject(storePayloadDic, (payloadValue) =>
-            modifyFn(payloadValue, docId),
+          storePayloadDic = mapObject(storePayloadDic, (payloadValue, storeName) =>
+            unwrapStoreSplits(modifyFn(payloadValue, docId), `${storeName}`),
           )
         }
 
diff --git a/packages/plugin-firestore/test/internal/storeSplit.test.ts b/packages/plugin-firestore/test/internal/storeSplit.test.ts
index 5d070f2e..5e4d78b1 100644
--- a/packages/plugin-firestore/test/internal/storeSplit.test.ts
+++ b/packages/plugin-firestore/test/internal/storeSplit.test.ts
@@ -37,3 +37,43 @@ import { firestoreDeepEqual } from '../helpers/firestoreDeepEqual.js'
     assert.deepEqual(doc.data?.base.HP as any, '9000')
   })
 }
+{
+  const testName = 'write: merge (document) with storeSplit in hook'
+  test(testName, async () => {
+    const { pokedexModule } = await createMagnetarInstance(testName, {
+      insertDocs: { 'pokedex/1': pokedex(1) },
+    })
+
+    const doc = pokedexModule.doc('1')
+    assert.deepEqual(doc.data, pokedex(1))
+    await firestoreDeepEqual(testName, 'pokedex/1', pokedex(1))
+
+    try {
+      await doc.merge(
+        { base: { HP: 9000 } },
+        {
+          syncDebounceMs: 1,
+          modifyPayloadOn: {
+            write: (payload) => {
+              assert.deepEqual(payload, { base: { HP: 9000 } })
+              return { base: { HP: storeSplit({ cache: 9000, remote: '9000' }) } }
+            },
+          },
+        },
+      )
+    } catch (error) {
+      assert.fail(JSON.stringify(error))
+    }
+
+    const mergedResult = merge(pokedex(1), { base: { HP: 9000 } })
+
+    assert.deepEqual(pokedexModule.data.get('1'), mergedResult)
+    assert.deepEqual(doc.data, mergedResult)
+
+    await firestoreDeepEqual(testName, 'pokedex/1', merge(pokedex(1), { base: { HP: '9000' } }))
+
+    const fetchedDoc = await doc.fetch({ force: true })
+    assert.deepEqual(fetchedDoc?.base.HP as any, '9000')
+    assert.deepEqual(doc.data?.base.HP as any, '9000')
+  })
+}