Fix broken extractRehydrationInfo

Author: markerikson

packages/toolkit/src/query/core/buildSlice.ts

@@ -519,13 +519,12 @@ export function buildSlice({
               providedTags as FullTagDescription<string>[]
           }
         },
-        prepare:
-          prepareAutoBatched<
-            Array<{
-              queryCacheKey: QueryCacheKey
-              providedTags: readonly FullTagDescription<string>[]
-            }>
-          >(),
+        prepare: prepareAutoBatched<
+          Array<{
+            queryCacheKey: QueryCacheKey
+            providedTags: readonly FullTagDescription<string>[]
+          }>
+        >(),
       },
     },
     extraReducers(builder) {
@@ -538,7 +537,9 @@ export function buildSlice({
         )
         .addMatcher(hasRehydrationInfo, (draft, action) => {
           const { provided } = extractRehydrationInfo(action)!
-          for (const [type, incomingTags] of Object.entries(provided)) {
+          for (const [type, incomingTags] of Object.entries(
+            provided.tags ?? {},
+          )) {
             for (const [id, cacheKeys] of Object.entries(incomingTags)) {
               const subscribedQueries = ((draft.tags[type] ??= {})[
                 id || '__internal_without_id'
@@ -549,6 +550,7 @@ export function buildSlice({
                 if (!alreadySubscribed) {
                   subscribedQueries.push(queryCacheKey)
                 }
+                draft.keys[queryCacheKey] = provided.keys[queryCacheKey]
               }
             }
           }

packages/toolkit/src/query/tests/buildSlice.test.ts

@@ -1,10 +1,15 @@
-import { createSlice } from '@reduxjs/toolkit'
+import { createSlice, createAction } from '@reduxjs/toolkit'
+import type { CombinedState } from '@reduxjs/toolkit/query'
 import { createApi } from '@reduxjs/toolkit/query'
 import { delay } from 'msw'
 import { setupApiStore } from '../../tests/utils/helpers'
 
 let shouldApiResponseSuccess = true
 
+const rehydrateAction = createAction<{ api: CombinedState<any, any, any> }>(
+  'persist/REHYDRATE',
+)
+
 const baseQuery = (args?: any) => ({ data: args })
 const api = createApi({
   baseQuery,
@@ -17,6 +22,12 @@ const api = createApi({
       providesTags: (result) => (result?.success ? ['SUCCEED'] : ['FAILED']),
     }),
   }),
+  extractRehydrationInfo(action, { reducerPath }) {
+    if (rehydrateAction.match(action)) {
+      return action.payload?.[reducerPath]
+    }
+    return undefined
+  },
 })
 const { getUser } = api.endpoints
 
@@ -114,6 +125,20 @@ describe('buildSlice', () => {
       api.util.selectInvalidatedBy(storeRef.store.getState(), ['FAILED']),
     ).toHaveLength(1)
   })
+
+  it('handles extractRehydrationInfo correctly', async () => {
+    await storeRef.store.dispatch(getUser.initiate(1))
+    await storeRef.store.dispatch(getUser.initiate(2))
+
+    const stateWithUser = storeRef.store.getState()
+
+    storeRef.store.dispatch(api.util.resetApiState())
+
+    storeRef.store.dispatch(rehydrateAction({ api: stateWithUser.api }))
+
+    const rehydratedState = storeRef.store.getState()
+    expect(rehydratedState).toEqual(stateWithUser)
+  })
 })
 
 describe('`merge` callback', () => {