fix(middleware/persist): ensure argument for onRehydrateStorage and onHydrate is defined on first hydration (#1692)

* add tests ensuring that onRehydrateStorage is always passed the latest state

* fix persist to always pass latest state to onRehydrateStorage and onHydrate listeners

* document that tests for onHydrate during first hydration are not possible

* ensure state updates during onHydrate are reflected in onRehydrateStorage callback

* undo modifications to persist's old implementation

src/middleware/persist.ts

@@ -421,11 +421,16 @@ const newImpl: PersistImpl = (config, baseOptions) => (set, get, api) => {
   const hydrate = () => {
     if (!storage) return
 
+    // On the first invocation of 'hydrate', state will not yet be defined (this is
+    // true for both the 'asynchronous' and 'synchronous' case). Pass 'configResult'
+    // as a backup  to 'get()' so listeners and 'onRehydrateStorage' are called with
+    // the latest available state.
+
     hasHydrated = false
-    hydrationListeners.forEach((cb) => cb(get()))
+    hydrationListeners.forEach((cb) => cb(get() ?? configResult))
 
     const postRehydrationCallback =
-      options.onRehydrateStorage?.(get()) || undefined
+      options.onRehydrateStorage?.(get() ?? configResult) || undefined
 
     // bind is used to avoid `TypeError: Illegal invocation` error
     return toThenable(storage.getItem.bind(storage))(options.name)

tests/persistAsync.test.tsx

@@ -385,6 +385,56 @@ describe('persist middleware with async configuration', () => {
     })
   })
 
+  it('passes the latest state to onRehydrateStorage and onHydrate on first hydrate', async () => {
+    const onRehydrateStorageSpy =
+      jest.fn<<S>(s: S) => (s?: S, e?: unknown) => void>()
+
+    const storage = {
+      getItem: async () => JSON.stringify({ state: { count: 1 } }),
+      setItem: () => {},
+      removeItem: () => {},
+    }
+
+    const useBoundStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        storage: createJSONStorage(() => storage),
+        onRehydrateStorage: onRehydrateStorageSpy,
+      })
+    )
+
+    /**
+     * NOTE: It's currently not possible to add an 'onHydrate' listener which will be
+     * invoked prior to the first hydration. This is because, during first hydration,
+     * the 'onHydrate' listener set (which will be empty) is evaluated before the
+     * 'persist' API is exposed to the caller of 'create'/'createStore'.
+     *
+     * const onHydrateSpy = jest.fn()
+     * useBoundStore.persist.onHydrate(onHydrateSpy)
+     * ...
+     * await waitFor(() => expect(onHydrateSpy).toBeCalledWith({ count: 0 }))
+     */
+
+    function Counter() {
+      const { count } = useBoundStore()
+      return <div>count: {count}</div>
+    }
+
+    const { findByText } = render(
+      <StrictMode>
+        <Counter />
+      </StrictMode>
+    )
+
+    await findByText('count: 1')
+
+    // The 'onRehydrateStorage' spy is invoked prior to rehydration, so it should
+    // be passed the default state.
+    await waitFor(() => {
+      expect(onRehydrateStorageSpy).toBeCalledWith({ count: 0 })
+    })
+  })
+
   it('gives the merged state to onRehydrateStorage', async () => {
     const onRehydrateStorageSpy = jest.fn()
 

tests/persistSync.test.tsx

@@ -226,6 +226,41 @@ describe('persist middleware with sync configuration', () => {
     )
   })
 
+  it('passes the latest state to onRehydrateStorage and onHydrate on first hydrate', () => {
+    const onRehydrateStorageSpy =
+      jest.fn<<S>(s: S) => (s?: S, e?: unknown) => void>()
+
+    const storage = {
+      getItem: () => JSON.stringify({ state: { count: 1 } }),
+      setItem: () => {},
+      removeItem: () => {},
+    }
+
+    const useBoundStore = create(
+      persist(() => ({ count: 0 }), {
+        name: 'test-storage',
+        storage: createJSONStorage(() => storage),
+        onRehydrateStorage: onRehydrateStorageSpy,
+      })
+    )
+
+    /**
+     * NOTE: It's currently not possible to add an 'onHydrate' listener which will be
+     * invoked prior to the first hydration. This is because, during first hydration,
+     * the 'onHydrate' listener set (which will be empty) is evaluated before the
+     * 'persist' API is exposed to the caller of 'create'/'createStore'.
+     *
+     * const onHydrateSpy = jest.fn()
+     * useBoundStore.persist.onHydrate(onHydrateSpy)
+     * expect(onHydrateSpy).toBeCalledWith({ count: 0 })
+     */
+
+    // The 'onRehydrateStorage' and 'onHydrate' spies are invoked prior to rehydration,
+    // so they should both be passed the default state.
+    expect(onRehydrateStorageSpy).toBeCalledWith({ count: 0 })
+    expect(useBoundStore.getState()).toEqual({ count: 1 })
+  })
+
   it('gives the merged state to onRehydrateStorage', () => {
     const onRehydrateStorageSpy = jest.fn()