From 9bae06922a361d525757cb277fba576caae05524 Mon Sep 17 00:00:00 2001
From: Danilo Britto <dbritto.dev@gmail.com>
Date: Wed, 4 Dec 2024 08:06:36 -0500
Subject: [PATCH] fix(persist): fix async migrate on persist middleware (#2877)

* fix(perist): add support for async migrate on persist middleware

* feat(vitest): rename vitest.config.ts to vitest.config.mts

* fix: fix minor issues

* fix: fix minor issues

* feat: minor changes

* feat: revert vitest changes

---------

Co-authored-by: Daishi Kato <dai-shi@users.noreply.github.com>
---
 src/middleware/persist.ts   | 15 ++++++++-------
 tests/persistAsync.test.tsx |  4 ++--
 tests/persistSync.test.tsx  |  2 +-
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/middleware/persist.ts b/src/middleware/persist.ts
index 6ec34d6c7f..58b2edd1b9 100644
--- a/src/middleware/persist.ts
+++ b/src/middleware/persist.ts
@@ -258,13 +258,14 @@ const persistImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
             deserializedStorageValue.version !== options.version
           ) {
             if (options.migrate) {
-              return [
-                true,
-                options.migrate(
-                  deserializedStorageValue.state,
-                  deserializedStorageValue.version,
-                ),
-              ] as const
+              const migration = options.migrate(
+                deserializedStorageValue.state,
+                deserializedStorageValue.version,
+              )
+              if (migration instanceof Promise) {
+                return migration.then((result) => [true, result] as const)
+              }
+              return [true, migration] as const
             }
             console.error(
               `State loaded from storage couldn't be migrated since no migrate function was provided`,
diff --git a/tests/persistAsync.test.tsx b/tests/persistAsync.test.tsx
index f86f3028cf..e36211e0ad 100644
--- a/tests/persistAsync.test.tsx
+++ b/tests/persistAsync.test.tsx
@@ -193,10 +193,10 @@ describe('persist middleware with async configuration', () => {
     })
   })
 
-  it('can migrate persisted state', async () => {
+  it('can async migrate persisted state', async () => {
     const setItemSpy = vi.fn()
     const onRehydrateStorageSpy = vi.fn()
-    const migrateSpy = vi.fn(() => ({ count: 99 }))
+    const migrateSpy = vi.fn(() => Promise.resolve({ count: 99 }))
 
     const storage = {
       getItem: async () =>
diff --git a/tests/persistSync.test.tsx b/tests/persistSync.test.tsx
index 5013bd6560..e4ad8c662d 100644
--- a/tests/persistSync.test.tsx
+++ b/tests/persistSync.test.tsx
@@ -131,7 +131,7 @@ describe('persist middleware with sync configuration', () => {
     expect(onRehydrateStorageSpy2).toBeCalledWith({ count: 42 }, undefined)
   })
 
-  it('can migrate persisted state', () => {
+  it('can non-async migrate persisted state', () => {
     const setItemSpy = vi.fn()
     const onRehydrateStorageSpy = vi.fn()
     const migrateSpy = vi.fn(() => ({ count: 99 }))