add infinite query support for updateCachedData

packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts

@@ -1,7 +1,7 @@
 import type { ThunkDispatch, UnknownAction } from '@reduxjs/toolkit'
 import type { BaseQueryFn, BaseQueryMeta } from '../../baseQueryTypes'
 import type { BaseEndpointDefinition } from '../../endpointDefinitions'
-import { DefinitionType } from '../../endpointDefinitions'
+import { DefinitionType, isAnyQueryDefinition } from '../../endpointDefinitions'
 import type { QueryCacheKey, RootState } from '../apiState'
 import type {
   MutationResultSelectorResult,
@@ -328,9 +328,7 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
     cacheDataLoaded.catch(() => {})
     lifecycleMap[queryCacheKey] = lifecycle
     const selector = (api.endpoints[endpointName] as any).select(
-      endpointDefinition.type === DefinitionType.query
-        ? originalArgs
-        : queryCacheKey,
+      isAnyQueryDefinition(endpointDefinition) ? originalArgs : queryCacheKey,
     )
 
     const extra = mwApi.dispatch((_, __, extra) => extra)
@@ -339,7 +337,7 @@ export const buildCacheLifecycleHandler: InternalHandlerBuilder = ({
       getCacheEntry: () => selector(mwApi.getState()),
       requestId,
       extra,
-      updateCachedData: (endpointDefinition.type === DefinitionType.query
+      updateCachedData: (isAnyQueryDefinition(endpointDefinition)
         ? (updateRecipe: Recipe<any>) =>
             mwApi.dispatch(
               api.util.updateQueryData(

packages/toolkit/src/query/core/buildMiddleware/queryLifecycle.ts

@@ -3,7 +3,7 @@ import type {
   BaseQueryFn,
   BaseQueryMeta,
 } from '../../baseQueryTypes'
-import { DefinitionType } from '../../endpointDefinitions'
+import { DefinitionType, isAnyQueryDefinition } from '../../endpointDefinitions'
 import type { Recipe } from '../buildThunks'
 import { isFulfilled, isPending, isRejected } from '../rtkImports'
 import type {
@@ -459,9 +459,7 @@ export const buildQueryLifecycleHandler: InternalHandlerBuilder = ({
         queryFulfilled.catch(() => {})
         lifecycleMap[requestId] = lifecycle
         const selector = (api.endpoints[endpointName] as any).select(
-          endpointDefinition.type === DefinitionType.query
-            ? originalArgs
-            : requestId,
+          isAnyQueryDefinition(endpointDefinition) ? originalArgs : requestId,
         )
 
         const extra = mwApi.dispatch((_, __, extra) => extra)
@@ -470,7 +468,7 @@ export const buildQueryLifecycleHandler: InternalHandlerBuilder = ({
           getCacheEntry: () => selector(mwApi.getState()),
           requestId,
           extra,
-          updateCachedData: (endpointDefinition.type === DefinitionType.query
+          updateCachedData: (isAnyQueryDefinition(endpointDefinition)
             ? (updateRecipe: Recipe<any>) =>
                 mwApi.dispatch(
                   api.util.updateQueryData(

packages/toolkit/src/query/endpointDefinitions.ts

@@ -912,6 +912,14 @@ export function isInfiniteQueryDefinition(
   return e.type === DefinitionType.infinitequery
 }
 
+export function isAnyQueryDefinition(
+  e: EndpointDefinition<any, any, any, any>,
+): e is
+  | QueryDefinition<any, any, any, any>
+  | InfiniteQueryDefinition<any, any, any, any, any> {
+  return isQueryDefinition(e) || isInfiniteQueryDefinition(e)
+}
+
 export type EndpointBuilder<
   BaseQuery extends BaseQueryFn,
   TagTypes extends string,

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

@@ -480,7 +480,7 @@ test(`mutation: getCacheEntry`, async () => {
   })
 })
 
-test('updateCachedData', async () => {
+test('query: updateCachedData', async () => {
   const trackCalls = vi.fn()
 
   const extended = api.injectEndpoints({
@@ -551,6 +551,95 @@ test('updateCachedData', async () => {
   })
 })
 
+test('updateCachedData - infinite query', async () => {
+  const trackCalls = vi.fn()
+
+  const extended = api.injectEndpoints({
+    overrideExisting: true,
+    endpoints: (build) => ({
+      infiniteInjected: build.infiniteQuery<{ value: string }, string, number>({
+        query: () => '/success',
+        infiniteQueryOptions: {
+          initialPageParam: 1,
+          getNextPageParam: (
+            lastPage,
+            allPages,
+            lastPageParam,
+            allPageParams,
+          ) => lastPageParam + 1,
+        },
+        async onCacheEntryAdded(
+          arg,
+          {
+            dispatch,
+            getState,
+            getCacheEntry,
+            updateCachedData,
+            cacheEntryRemoved,
+            cacheDataLoaded,
+          },
+        ) {
+          expect(getCacheEntry().data).toEqual(undefined)
+          // calling `updateCachedData` when there is no data yet should not do anything
+          updateCachedData((draft) => {
+            draft.pages = [{ value: 'TEST' }]
+            draft.pageParams = [1]
+            trackCalls()
+          })
+          expect(trackCalls).not.toHaveBeenCalled()
+          expect(getCacheEntry().data).toEqual(undefined)
+
+          gotFirstValue(await cacheDataLoaded)
+
+          expect(getCacheEntry().data).toEqual({
+            pages: [{ value: 'success' }],
+            pageParams: [1],
+          })
+          updateCachedData((draft) => {
+            draft.pages = [{ value: 'TEST' }]
+            draft.pageParams = [1]
+            trackCalls()
+          })
+          expect(trackCalls).toHaveBeenCalledOnce()
+          expect(getCacheEntry().data).toEqual({
+            pages: [{ value: 'TEST' }],
+            pageParams: [1],
+          })
+
+          await cacheEntryRemoved
+
+          expect(getCacheEntry().data).toEqual(undefined)
+          // calling `updateCachedData` when there is no data any more should not do anything
+          updateCachedData((draft) => {
+            draft.pages = [{ value: 'TEST' }, { value: 'TEST2' }]
+            draft.pageParams = [1, 2]
+            trackCalls()
+          })
+          expect(trackCalls).toHaveBeenCalledOnce()
+          expect(getCacheEntry().data).toEqual(undefined)
+
+          onCleanup()
+        },
+      }),
+    }),
+  })
+  const promise = storeRef.store.dispatch(
+    extended.endpoints.infiniteInjected.initiate('arg'),
+  )
+  await promise
+  promise.unsubscribe()
+
+  await fakeTimerWaitFor(() => {
+    expect(gotFirstValue).toHaveBeenCalled()
+  })
+
+  await vi.advanceTimersByTimeAsync(61000)
+
+  await fakeTimerWaitFor(() => {
+    expect(onCleanup).toHaveBeenCalled()
+  })
+})
+
 test('dispatching further actions does not trigger another lifecycle', async () => {
   const extended = api.injectEndpoints({
     overrideExisting: true,

packages/toolkit/src/query/tests/queryLifecycle.test.tsx

@@ -342,8 +342,6 @@ test('mutation: getCacheEntry (error)', async () => {
 })
 
 test('query: updateCachedData', async () => {
-  const trackCalls = vi.fn()
-
   const extended = api.injectEndpoints({
     overrideExisting: true,
     endpoints: (build) => ({
@@ -366,7 +364,7 @@ test('query: updateCachedData', async () => {
           })
 
           try {
-            const val = await queryFulfilled
+            await queryFulfilled
             onSuccess(getCacheEntry().data)
           } catch (error) {
             updateCachedData((draft) => {
@@ -423,6 +421,106 @@ test('query: updateCachedData', async () => {
   onSuccess.mockClear()
 })
 
+test('infinite query: updateCachedData', async () => {
+  const extended = api.injectEndpoints({
+    overrideExisting: true,
+    endpoints: (build) => ({
+      infiniteInjected: build.infiniteQuery<{ value: string }, string, number>({
+        query: () => '/success',
+        infiniteQueryOptions: {
+          initialPageParam: 1,
+          getNextPageParam: (
+            lastPage,
+            allPages,
+            lastPageParam,
+            allPageParams,
+          ) => lastPageParam + 1,
+        },
+        async onQueryStarted(
+          arg,
+          {
+            dispatch,
+            getState,
+            getCacheEntry,
+            updateCachedData,
+            queryFulfilled,
+          },
+        ) {
+          // calling `updateCachedData` when there is no data yet should not do anything
+          // but if there is a cache value it will be updated & overwritten by the next successful result
+          updateCachedData((draft) => {
+            draft.pages = [{ value: '.' }]
+            draft.pageParams = [1]
+          })
+
+          try {
+            await queryFulfilled
+            onSuccess(getCacheEntry().data)
+          } catch (error) {
+            updateCachedData((draft) => {
+              draft.pages = [{ value: 'success.x' }]
+              draft.pageParams = [1]
+            })
+            onError(getCacheEntry().data)
+          }
+        },
+      }),
+    }),
+  })
+
+  // request 1: success
+  expect(onSuccess).not.toHaveBeenCalled()
+  storeRef.store.dispatch(extended.endpoints.infiniteInjected.initiate('arg'))
+
+  await waitFor(() => {
+    expect(onSuccess).toHaveBeenCalled()
+  })
+  expect(onSuccess).toHaveBeenCalledWith({
+    pages: [{ value: 'success' }],
+    pageParams: [1],
+  })
+  onSuccess.mockClear()
+
+  // request 2: error
+  expect(onError).not.toHaveBeenCalled()
+  server.use(
+    http.get(
+      'https://example.com/success',
+      () => {
+        return HttpResponse.json({ value: 'failed' }, { status: 500 })
+      },
+      { once: true },
+    ),
+  )
+  storeRef.store.dispatch(
+    extended.endpoints.infiniteInjected.initiate('arg', { forceRefetch: true }),
+  )
+
+  await waitFor(() => {
+    expect(onError).toHaveBeenCalled()
+  })
+  expect(onError).toHaveBeenCalledWith({
+    pages: [{ value: 'success.x' }],
+    pageParams: [1],
+  })
+
+  // request 3: success
+  expect(onSuccess).not.toHaveBeenCalled()
+
+  storeRef.store.dispatch(
+    extended.endpoints.infiniteInjected.initiate('arg', { forceRefetch: true }),
+  )
+
+  await waitFor(() => {
+    expect(onSuccess).toHaveBeenCalled()
+  })
+  expect(onSuccess).toHaveBeenCalledWith({
+    pages: [{ value: 'success' }],
+    pageParams: [1],
+  })
+  onSuccess.mockClear()
+})
+
 test('query: will only start lifecycle if query is not skipped due to `condition`', async () => {
   const extended = api.injectEndpoints({
     overrideExisting: true,