Fix upsertQueryEntries for infinite queries

Author: markerikson

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

@@ -29,6 +29,9 @@ import type {
 } from './apiState'
 import { QueryStatus } from './apiState'
 import type {
+  AllQueryKeys,
+  QueryArgFromAnyQueryDefinition,
+  DataFromAnyQueryDefinition,
   InfiniteQueryThunk,
   MutationThunk,
   QueryThunk,
@@ -64,11 +67,11 @@ import type { InternalSerializeQueryArgs } from '../defaultSerializeQueryArgs'
  */
 export type NormalizedQueryUpsertEntry<
   Definitions extends EndpointDefinitions,
-  EndpointName extends QueryKeys<Definitions>,
+  EndpointName extends AllQueryKeys<Definitions>,
 > = {
   endpointName: EndpointName
-  arg: QueryArgFrom<Definitions[EndpointName]>
-  value: ResultTypeFrom<Definitions[EndpointName]>
+  arg: QueryArgFromAnyQueryDefinition<Definitions, EndpointName>
+  value: DataFromAnyQueryDefinition<Definitions, EndpointName>
 }
 
 /**
@@ -89,7 +92,7 @@ export type ProcessedQueryUpsertEntry = {
  * A typesafe representation of a util action creator that accepts cache entry descriptions to upsert
  */
 export type UpsertEntries<Definitions extends EndpointDefinitions> = (<
-  EndpointNames extends Array<QueryKeys<Definitions>>,
+  EndpointNames extends Array<AllQueryKeys<Definitions>>,
 >(
   entries: [
     ...{

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

@@ -6,6 +6,7 @@ import type {
   ThunkDispatch,
   UnknownAction,
 } from '@reduxjs/toolkit'
+import util from 'util'
 import type { Patch } from 'immer'
 import { isDraftable, produceWithPatches } from 'immer'
 import type { Api, ApiContext } from '../apiTypes'
@@ -198,11 +199,11 @@ export type PatchQueryDataThunk<
   updateProvided?: boolean,
 ) => ThunkAction<void, PartialState, any, UnknownAction>
 
-type AllQueryKeys<Definitions extends EndpointDefinitions> =
+export type AllQueryKeys<Definitions extends EndpointDefinitions> =
   | QueryKeys<Definitions>
   | InfiniteQueryKeys<Definitions>
 
-type QueryArgFromAnyQueryDefinition<
+export type QueryArgFromAnyQueryDefinition<
   Definitions extends EndpointDefinitions,
   EndpointName extends AllQueryKeys<Definitions>,
 > =
@@ -218,7 +219,7 @@ type QueryArgFromAnyQueryDefinition<
       ? QueryArgFrom<Definitions[EndpointName]>
       : never
 
-type DataFromAnyQueryDefinition<
+export type DataFromAnyQueryDefinition<
   Definitions extends EndpointDefinitions,
   EndpointName extends AllQueryKeys<Definitions>,
 > =
@@ -237,7 +238,7 @@ type DataFromAnyQueryDefinition<
       ? ResultTypeFrom<Definitions[EndpointName]>
       : unknown
 
-type UpsertThunkResult<
+export type UpsertThunkResult<
   Definitions extends EndpointDefinitions,
   EndpointName extends AllQueryKeys<Definitions>,
 > =

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

@@ -557,9 +557,8 @@ describe('Infinite queries', () => {
       }),
     )
 
-    const entry1Updated = pokemonApi.endpoints.getInfinitePokemon.select(
-      'fire',
-    )(storeRef.store.getState())
+    const selectFire = pokemonApi.endpoints.getInfinitePokemon.select('fire')
+    const entry1Updated = selectFire(storeRef.store.getState())
 
     expect(entry1Updated.data).toEqual({
       pages: [
@@ -577,18 +576,49 @@ describe('Infinite queries', () => {
     )
 
     const entry2InitialLoad = await res2
-    const entry2Updated = pokemonApi.endpoints.getInfinitePokemon.select(
-      'water',
-    )(storeRef.store.getState())
-
-    console.log(
-      'entry2Updated',
-      util.inspect(entry2Updated, { depth: Infinity }),
-    )
+    const selectWater = pokemonApi.endpoints.getInfinitePokemon.select('water')
+    const entry2Updated = selectWater(storeRef.store.getState())
 
     expect(entry2Updated.data).toEqual({
       pages: [[{ id: '2', name: 'Pokemon 2' }]],
       pageParams: [2],
     })
+
+    storeRef.store.dispatch(
+      pokemonApi.util.upsertQueryEntries([
+        {
+          endpointName: 'getInfinitePokemon',
+          arg: 'air',
+          value: {
+            pages: [[{ id: '3', name: 'Pokemon 3' }]],
+            pageParams: [3],
+          },
+        },
+      ]),
+    )
+
+    const selectAir = pokemonApi.endpoints.getInfinitePokemon.select('air')
+    const entry3Initial = selectAir(storeRef.store.getState())
+
+    expect(entry3Initial.data).toEqual({
+      pages: [[{ id: '3', name: 'Pokemon 3' }]],
+      pageParams: [3],
+    })
+
+    await storeRef.store.dispatch(
+      pokemonApi.endpoints.getInfinitePokemon.initiate('air', {
+        direction: 'forward',
+      }),
+    )
+
+    const entry3Updated = selectAir(storeRef.store.getState())
+
+    expect(entry3Updated.data).toEqual({
+      pages: [
+        [{ id: '3', name: 'Pokemon 3' }],
+        [{ id: '4', name: 'Pokemon 4' }],
+      ],
+      pageParams: [3, 4],
+    })
   })
 })