From e101ec6a281d783c3692f52f86dc73baafc12342 Mon Sep 17 00:00:00 2001
From: Mark Erikson <mark@isquaredsoftware.com>
Date: Fri, 13 Dec 2024 21:13:05 -0500
Subject: [PATCH] Ensure upserted cache entries always get written

---
 packages/toolkit/src/query/core/buildSlice.ts | 19 ++--
 .../query/tests/optimisticUpserts.test.tsx    | 89 ++++++++++++++++++-
 2 files changed, 100 insertions(+), 8 deletions(-)

diff --git a/packages/toolkit/src/query/core/buildSlice.ts b/packages/toolkit/src/query/core/buildSlice.ts
index 11e0c863c5..0c5a42becb 100644
--- a/packages/toolkit/src/query/core/buildSlice.ts
+++ b/packages/toolkit/src/query/core/buildSlice.ts
@@ -84,7 +84,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> = <
+export type UpsertEntries<Definitions extends EndpointDefinitions> = (<
   EndpointNames extends Array<QueryKeys<Definitions>>,
 >(
   entries: [
@@ -95,7 +95,11 @@ export type UpsertEntries<Definitions extends EndpointDefinitions> = <
       >
     },
   ],
-) => PayloadAction<NormalizedQueryUpsertEntryPayload[]>
+) => PayloadAction<NormalizedQueryUpsertEntryPayload[]>) & {
+  match: (
+    action: unknown,
+  ) => action is PayloadAction<NormalizedQueryUpsertEntryPayload[]>
+}
 
 function updateQuerySubstateIfExists(
   state: QueryState<any>,
@@ -212,10 +216,10 @@ export function buildSlice({
       // RTK_autoBatch: true
     },
     payload: unknown,
+    upserting: boolean,
   ) {
     updateQuerySubstateIfExists(draft, meta.arg.queryCacheKey, (substate) => {
-      if (substate.requestId !== meta.requestId && !isUpsertQuery(meta.arg))
-        return
+      if (substate.requestId !== meta.requestId && !upserting) return
       const { merge } = definitions[meta.arg.endpointName] as QueryDefinition<
         any,
         any,
@@ -248,7 +252,7 @@ export function buildSlice({
       } else {
         // Assign or safely update the cache data.
         substate.data =
-          definitions[meta.arg.endpointName].structuralSharing ?? true
+          (definitions[meta.arg.endpointName].structuralSharing ?? true)
             ? copyWithStructuralSharing(
                 isDraft(substate.data)
                   ? original(substate.data)
@@ -307,6 +311,8 @@ export function buildSlice({
                 baseQueryMeta: {},
               },
               value,
+              // We know we're upserting here
+              true,
             )
           }
         },
@@ -365,7 +371,8 @@ export function buildSlice({
           writePendingCacheEntry(draft, arg, upserting, meta)
         })
         .addCase(queryThunk.fulfilled, (draft, { meta, payload }) => {
-          writeFulfilledCacheEntry(draft, meta, payload)
+          const upserting = isUpsertQuery(meta.arg)
+          writeFulfilledCacheEntry(draft, meta, payload, upserting)
         })
         .addCase(
           queryThunk.rejected,
diff --git a/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx b/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx
index ac6c8ee473..d7099c1c03 100644
--- a/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx
+++ b/packages/toolkit/src/query/tests/optimisticUpserts.test.tsx
@@ -5,7 +5,13 @@ import {
   hookWaitFor,
   setupApiStore,
 } from '../../tests/utils/helpers'
-import { renderHook, act, waitFor } from '@testing-library/react'
+import {
+  render,
+  renderHook,
+  act,
+  waitFor,
+  screen,
+} from '@testing-library/react'
 import { delay } from 'msw'
 
 interface Post {
@@ -14,6 +20,11 @@ interface Post {
   contents: string
 }
 
+interface FolderT {
+  id: number
+  children: FolderT[]
+}
+
 const baseQuery = vi.fn()
 beforeEach(() => baseQuery.mockReset())
 
@@ -28,7 +39,7 @@ const api = createApi({
         .catch((e: any) => ({ error: e }))
     return { data: result, meta: 'meta' }
   },
-  tagTypes: ['Post'],
+  tagTypes: ['Post', 'Folder'],
   endpoints: (build) => ({
     getPosts: build.query<Post[], void>({
       query: () => '/posts',
@@ -80,6 +91,30 @@ const api = createApi({
       },
       keepUnusedDataFor: 0.01,
     }),
+    getFolder: build.query<FolderT, number>({
+      queryFn: async (args) => {
+        return {
+          data: {
+            id: args,
+            // Folder contains children that are as well folders
+            children: [{ id: 2, children: [] }],
+          },
+        }
+      },
+      providesTags: (result, err, args) => [{ type: 'Folder', id: args }],
+      onQueryStarted: async (args, queryApi) => {
+        const { data } = await queryApi.queryFulfilled
+
+        // Upsert getFolder endpoint with children from response data
+        const upsertData = data.children.map((child) => ({
+          arg: child.id,
+          endpointName: 'getFolder' as const,
+          value: child,
+        }))
+
+        queryApi.dispatch(api.util.upsertQueryEntries(upsertData))
+      },
+    }),
   }),
 })
 
@@ -434,6 +469,56 @@ describe('upsertQueryEntries', () => {
       undefined,
     )
   })
+
+  test('Handles repeated upserts and async lifecycles', async () => {
+    const StateForUpsertFolder = ({ folderId }: { folderId: number }) => {
+      const { status } = api.useGetFolderQuery(folderId)
+
+      return (
+        <>
+          <div>
+            Status getFolder with ID (
+            {folderId === 1 ? 'original request' : 'upserted'}) {folderId}:{' '}
+            <span data-testid={`status-${folderId}`}>{status}</span>
+          </div>
+        </>
+      )
+    }
+
+    const Folder = () => {
+      const { data, isLoading, isError } = api.useGetFolderQuery(1)
+
+      return (
+        <div>
+          <h1>Folders</h1>
+
+          {isLoading && <div>Loading...</div>}
+
+          {isError && <div>Error...</div>}
+
+          <StateForUpsertFolder key={`state-${1}`} folderId={1} />
+          <StateForUpsertFolder key={`state-${2}`} folderId={2} />
+        </div>
+      )
+    }
+
+    render(<Folder />, {
+      wrapper: storeRef.wrapper,
+    })
+
+    await waitFor(() => {
+      const { actions } = storeRef.store.getState()
+      // Inspection:
+      // - 2 inits
+      // - 2 pendings, 2 fulfilleds for the hook queries
+      // - 2 upserts
+      expect(actions.length).toBe(8)
+      expect(
+        actions.filter((a) => api.util.upsertQueryEntries.match(a)).length,
+      ).toBe(2)
+    })
+    expect(screen.getByTestId('status-2').textContent).toBe('fulfilled')
+  })
 })
 
 describe('full integration', () => {