Fix bugs and add tests for sync storage

src/middleware.ts

@@ -274,60 +274,69 @@ export const persist = <S extends State>(
       }
     }
   }
-  // rehydrate initial state with existing stored state
+
+  let stateFromStorage
+    // rehydrate initial state with existing stored state
   ;(() => {
     const postRehydrationCallback = onRehydrateStorage?.(get()) || undefined
 
-    let storageValuePending = makeThenable(storage.getItem(name))
-    storageValuePending
-      .then((storageValue) => {
-        if (storageValue) {
-          return makeThenable(deserialize(storageValue))
-        } else {
-          return makeThenable(null)
-        }
-      })
-      .then((deserializedStorageValue) => {
-        if (deserializedStorageValue) {
-          if (deserializedStorageValue.version !== version) {
-            return makeThenable(
-              migrate?.(
+    try {
+      let storageValuePending = makeThenable(storage.getItem(name))
+      storageValuePending
+        .then((storageValue) => {
+          if (storageValue) {
+            return makeThenable(deserialize(storageValue))
+          } else {
+            return makeThenable(null)
+          }
+        })
+        .then((deserializedStorageValue) => {
+          if (deserializedStorageValue) {
+            if (deserializedStorageValue.version !== version) {
+              stateFromStorage = migrate?.(
                 deserializedStorageValue.state,
                 deserializedStorageValue.version
               )
-            )
+
+              return makeThenable(stateFromStorage)
+            } else {
+              stateFromStorage = deserializedStorageValue.state
+              set(stateFromStorage)
+              return makeThenable(null)
+            }
           } else {
-            set(deserializedStorageValue.state)
             return makeThenable(null)
           }
-        } else {
-          return makeThenable(null)
-        }
-      })
-      .then((migratedState) => {
-        if (migratedState) {
-          set(migratedState)
-          return makeThenable(setItem())
-        } else {
+        })
+        .then((migratedState) => {
+          if (migratedState) {
+            set(migratedState)
+            return makeThenable(setItem())
+          } else {
+            return makeThenable(null)
+          }
+        })
+        .then(() => {
+          postRehydrationCallback?.(get(), undefined)
           return makeThenable(null)
-        }
-      })
-      .then(() => {
-        postRehydrationCallback?.(get(), undefined)
-        return makeThenable(null)
-      })
-      // this handler will only catch promise related errors
-      .catch((e) => {
-        postRehydrationCallback?.(undefined, e)
-      })
+        })
+        // this handler will only catch promise related errors
+        .catch((e) => {
+          postRehydrationCallback?.(undefined, e)
+        })
+    } catch (e) {
+      postRehydrationCallback?.(undefined, e)
+    }
   })()
 
-  return config(
+  const configRes = config(
     (...args) => {
       set(...args)
       void setItem()
     },
     get,
     api
   )
+  const mergedRes = { ...configRes, ...(stateFromStorage || {}) }
+  return mergedRes
 }

tests/persistSync.test.tsx

@@ -0,0 +1,173 @@
+import create from '../src/index'
+import { persist } from '../src/middleware'
+
+const consoleError = console.error
+afterEach(() => {
+  console.error = consoleError
+})
+
+it('can rehydrate state', () => {
+  let postRehydrationCallbackCallCount = 0
+
+  const storage = {
+    getItem: (name: string) =>
+      JSON.stringify({
+        state: { count: 42, name },
+        version: 0,
+      }),
+    setItem: () => {},
+  }
+
+  const useStore = create(
+    persist(
+      () => ({
+        count: 0,
+        name: 'empty',
+      }),
+      {
+        name: 'test-storage',
+        getStorage: () => storage,
+        onRehydrateStorage: () => (state, error) => {
+          postRehydrationCallbackCallCount++
+          expect(error).toBeUndefined()
+          expect(state?.count).toBe(42)
+          expect(state?.name).toBe('test-storage')
+        },
+      }
+    )
+  )
+
+  expect(useStore.getState()).toEqual({ count: 42, name: 'test-storage' })
+  expect(postRehydrationCallbackCallCount).toBe(1)
+})
+
+it('can throw rehydrate error', () => {
+  let postRehydrationCallbackCallCount = 0
+
+  const storage = {
+    getItem: () => {
+      throw new Error('getItem error')
+    },
+    setItem: () => {},
+  }
+
+  create(
+    persist(() => ({ count: 0 }), {
+      name: 'test-storage',
+      getStorage: () => storage,
+      onRehydrateStorage: () => (_, e) => {
+        postRehydrationCallbackCallCount++
+        expect(e?.message).toBe('getItem error')
+      },
+    })
+  )
+
+  expect(postRehydrationCallbackCallCount).toBe(1)
+})
+
+it('can persist state', () => {
+  let setItemCallCount = 0
+
+  const storage = {
+    getItem: () => null,
+    setItem: (name: string, value: string) => {
+      setItemCallCount++
+      expect(name).toBe('test-storage')
+      expect(value).toBe(
+        JSON.stringify({
+          state: { count: 42 },
+          version: 0,
+        })
+      )
+    },
+  }
+
+  const useStore = create<any>(
+    persist(() => ({ count: 0 }), {
+      name: 'test-storage',
+      getStorage: () => storage,
+      onRehydrateStorage: () => (_, error) => {
+        expect(error).toBeUndefined()
+      },
+    })
+  )
+
+  expect(useStore.getState()).toEqual({ count: 0 })
+  useStore.setState({ count: 42 })
+  expect(useStore.getState()).toEqual({ count: 42 })
+  expect(setItemCallCount).toBe(1)
+})
+
+it('can migrate persisted state', () => {
+  let migrateCallCount = 0
+  let setItemCallCount = 0
+
+  const storage = {
+    getItem: () =>
+      JSON.stringify({
+        state: { count: 42 },
+        version: 12,
+      }),
+    setItem: (_: string, value: string) => {
+      setItemCallCount++
+      expect(value).toBe(
+        JSON.stringify({
+          state: { count: 99 },
+          version: 13,
+        })
+      )
+    },
+  }
+
+  const useStore = create(
+    persist(() => ({ count: 0 }), {
+      name: 'test-storage',
+      version: 13,
+      getStorage: () => storage,
+      onRehydrateStorage: () => (_, error) => {
+        expect(error).toBeUndefined()
+      },
+      migrate: (state, version) => {
+        migrateCallCount++
+        expect(state.count).toBe(42)
+        expect(version).toBe(12)
+        return { count: 99 }
+      },
+    })
+  )
+
+  expect(useStore.getState()).toEqual({ count: 99 })
+  expect(migrateCallCount).toBe(1)
+  expect(setItemCallCount).toBe(1)
+})
+
+it('can throw migrate error', () => {
+  let postRehydrationCallbackCallCount = 0
+
+  const storage = {
+    getItem: () =>
+      JSON.stringify({
+        state: {},
+        version: 12,
+      }),
+    setItem: () => {},
+  }
+
+  const useStore = create(
+    persist(() => ({ count: 0 }), {
+      name: 'test-storage',
+      version: 13,
+      getStorage: () => storage,
+      migrate: () => {
+        throw new Error('migrate error')
+      },
+      onRehydrateStorage: () => (_, e) => {
+        postRehydrationCallbackCallCount++
+        expect(e?.message).toBe('migrate error')
+      },
+    })
+  )
+
+  expect(useStore.getState()).toEqual({ count: 0 })
+  expect(postRehydrationCallbackCallCount).toBe(1)
+})