From 5d711981aab71ed31e07de21fb56a1254d33f95f Mon Sep 17 00:00:00 2001
From: Kilian Finger <finger@trafficon.eu>
Date: Mon, 4 Nov 2024 07:51:29 +0100
Subject: [PATCH] test: add async migrate test

---
 tests/persistAsync.test.tsx | 51 ++++++++++++++++++++++++++++++++++++-
 tests/persistSync.test.tsx  | 39 +++++++++++++++++++++++++++-
 2 files changed, 88 insertions(+), 2 deletions(-)

diff --git a/tests/persistAsync.test.tsx b/tests/persistAsync.test.tsx
index a7a5728090..8a4d421e2d 100644
--- a/tests/persistAsync.test.tsx
+++ b/tests/persistAsync.test.tsx
@@ -193,7 +193,7 @@ describe('persist middleware with async configuration', () => {
     })
   })
 
-  it('can migrate persisted state', async () => {
+  it('can non-async migrate persisted state', async () => {
     const setItemSpy = vi.fn()
     const onRehydrateStorageSpy = vi.fn()
     const migrateSpy = vi.fn(() => ({ count: 99 }))
@@ -242,6 +242,55 @@ describe('persist middleware with async configuration', () => {
     expect(onRehydrateStorageSpy).toBeCalledWith({ count: 99 }, undefined)
   })
 
+  it('can async migrate persisted state', async () => {
+    const setItemSpy = vi.fn()
+    const onRehydrateStorageSpy = vi.fn()
+    const migrateSpy = vi.fn(() => Promise.resolve({ count: 99 }))
+
+    const storage = {
+      getItem: async () =>
+        JSON.stringify({
+          state: { count: 42 },
+          version: 12,
+        }),
+      setItem: setItemSpy,
+      removeItem: () => {},
+    }
+
+    const useBoundStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        version: 13,
+        storage: createJSONStorage(() => storage),
+        onRehydrateStorage: () => onRehydrateStorageSpy,
+        migrate: migrateSpy,
+      }),
+    )
+
+    function Counter() {
+      const { count } = useBoundStore()
+      return <div>count: {count}</div>
+    }
+
+    render(
+      <StrictMode>
+        <Counter />
+      </StrictMode>,
+    )
+
+    await screen.findByText('count: 0')
+    await screen.findByText('count: 99')
+    expect(migrateSpy).toBeCalledWith({ count: 42 }, 12)
+    expect(setItemSpy).toBeCalledWith(
+      'test-storage',
+      JSON.stringify({
+        state: { count: 99 },
+        version: 13,
+      }),
+    )
+    expect(onRehydrateStorageSpy).toBeCalledWith({ count: 99 }, undefined)
+  })
+
   it('can merge partial persisted state', async () => {
     const storage = {
       getItem: async () =>
diff --git a/tests/persistSync.test.tsx b/tests/persistSync.test.tsx
index 5013bd6560..a450f0f529 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 }))
@@ -168,6 +168,43 @@ describe('persist middleware with sync configuration', () => {
     expect(onRehydrateStorageSpy).toBeCalledWith({ count: 99 }, undefined)
   })
 
+  it('can async migrate persisted state', () => {
+    const setItemSpy = vi.fn()
+    const onRehydrateStorageSpy = vi.fn()
+    const migrateSpy = vi.fn(() => Promise.resolve({ count: 99 }))
+
+    const storage = {
+      getItem: () =>
+        JSON.stringify({
+          state: { count: 42 },
+          version: 12,
+        }),
+      setItem: setItemSpy,
+      removeItem: () => {},
+    }
+
+    const useBoundStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        version: 13,
+        storage: createJSONStorage(() => storage),
+        onRehydrateStorage: () => onRehydrateStorageSpy,
+        migrate: migrateSpy,
+      }),
+    )
+
+    expect(useBoundStore.getState()).toEqual({ count: 99 })
+    expect(migrateSpy).toBeCalledWith({ count: 42 }, 12)
+    expect(setItemSpy).toBeCalledWith(
+      'test-storage',
+      JSON.stringify({
+        state: { count: 99 },
+        version: 13,
+      }),
+    )
+    expect(onRehydrateStorageSpy).toBeCalledWith({ count: 99 }, undefined)
+  })
+
   it('can correclty handle a missing migrate function', () => {
     console.error = vi.fn()
     const onRehydrateStorageSpy = vi.fn()