test: add async migrate test

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 () =>

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()