Await invalidations after mutations

app/api/__tests__/client.browser.spec.ts → app/api/__tests__/client.browser.spec.tsx

@@ -5,13 +5,16 @@
  *
  * Copyright Oxide Computer Company
  */
+import { QueryClientProvider, useQuery } from '@tanstack/react-query'
 import { http, HttpResponse } from 'msw'
 import { setupWorker } from 'msw/browser'
+import type { ReactNode } from 'react'
 import { afterAll, afterEach, beforeAll, describe, expect, it } from 'vitest'
+import { render } from 'vitest-browser-react'
 
 import { project } from '@oxide/api-mocks'
 
-import { api, q } from '..'
+import { api, type Project, q, queryClient, type ResultsPage, useApiMutation } from '..'
 import { resetDb } from '../../../mock-api/msw/db'
 import { handlers } from '../../../mock-api/msw/handlers'
 import { processServerError } from '../errors'
@@ -31,6 +34,7 @@
 afterEach(() => {
   resetDb()
   worker.resetHandlers()
+  queryClient.clear()
 })
 afterAll(() => worker.stop())
 
@@ -54,13 +58,35 @@
   )
 }
 
-// useApiQuery and useApiMutation are almost entirely typed wrappers around React
-// Query's useQuery and useMutation, so they're exercised end-to-end by the
-// Playwright suite (every error toast goes through this path). The logic worth
-// unit-testing directly is response parsing in the generated client
-// (`handleResponse`) and the error transformation in `processServerError` (the
-// latter is covered exhaustively in errors.spec.ts). These tests call the API
-// methods directly — no React, no renderHook — since they return an ApiResult.
+// The API hooks are mostly typed wrappers around React Query and are exercised
+// end-to-end by the Playwright suite. Most tests here therefore call API methods
+// directly; the mutation invalidation test renders a component because the
+// mutation's loading state is the behavior under test.
+
+const QueryClientWrapper = ({ children }: { children: ReactNode }) => (
+  <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+)
+
+function MutationInvalidationTest({ onSuccess }: { onSuccess: () => void }) {
+  const projects = useQuery(q(api.projectList, {}))
+  const createProject = useApiMutation(api.projectCreate, {
+    invalidateEndpoints: ['projectList'],
+    onSuccess,
+  })
+  const count = projects.data?.items.length ?? 0
+
+  return (
+    <button
+      type="button"
+      disabled={createProject.isPending}
+      onClick={() =>
+        createProject.mutate({ body: { name: 'new-project', description: '' } })
+      }
+    >
+      {createProject.isPending ? 'Creating' : 'Create'} project ({count} projects)
+    </button>
+  )
+}
 
 describe('API response parsing', () => {
   it('returns success data for a normal response', async () => {
@@ -126,3 +152,39 @@
   const queryOptions = q(api.siloView, params)
   expect(queryOptions.queryKey).toEqual(['siloView', params])
 })
+
+it('stays pending until invalidated queries have refreshed', async () => {
+  // capture the cached list length at onSuccess call time so we can assert
+  // onSuccess ran after the invalidated query refetched
+  let countAtSuccess: number | undefined
+  const onSuccess = () => {
+    const projects = queryClient.getQueryData<ResultsPage<Project>>(['projectList', {}])
+    countAtSuccess = projects?.items.length
+  }
+  const screen = await render(<MutationInvalidationTest onSuccess={onSuccess} />, {
+    wrapper: QueryClientWrapper,
+  })
+  const createButton = screen.getByRole('button')
+  await expect.element(createButton).toHaveAccessibleName('Create project (3 projects)')
+
+  const { promise: refetch, resolve: releaseRefetch } = Promise.withResolvers<void>()
+  worker.use(
+    http.get(
+      'http://testhost/v1/projects',
+      async () => {
+        await refetch
+        return HttpResponse.json({ items: [project, project, project, project] })
+      },
+      { once: true }
+    )
+  )
+
+  await createButton.click()
+  await expect.element(createButton).toBeDisabled()
+  await expect.element(createButton).toHaveAccessibleName('Creating project (3 projects)')
+
+  releaseRefetch()
+  await expect.element(createButton).toBeEnabled()
+  await expect.element(createButton).toHaveAccessibleName('Create project (4 projects)')
+  expect(countAtSuccess).toEqual(4)
+})

app/api/__tests__/safety.spec.ts

@@ -67,7 +67,7 @@ it('mock-api is only referenced in test files', () => {
   expect(grepFiles('api-mocks')).toMatchInlineSnapshot(`
     [
       "AGENTS.md",
-      "app/api/__tests__/client.browser.spec.ts",
+      "app/api/__tests__/client.browser.spec.tsx",
       "mock-api/msw/db.ts",
       "test/e2e/fleet-access.e2e.ts",
       "test/e2e/instance-create.e2e.ts",
@@ -83,7 +83,7 @@ it('mock-api is only referenced in test files', () => {
     [
       "AGENTS.md",
       "README.md",
-      "app/api/__tests__/client.browser.spec.ts",
+      "app/api/__tests__/client.browser.spec.tsx",
       "app/main.tsx",
       "app/msw-mock-api.ts",
       "docs/mock-api-differences.md",

app/api/client.ts

@@ -377,18 +377,45 @@ export const qErrorsAllowed = <Params, Data>(
 // object. You can only initalize the meta at the site of the useMutation call,
 // which doesn't work for the image upload use case because the timeout signal
 // needs to be initialized separately for each call.
+type ApiMutationOptions<Params, Data> = Omit<
+  UseMutationOptions<Data, ApiError, Params & { __signal?: AbortSignal }>,
+  'mutationFn' | 'onSettled'
+> & {
+  /**
+   * Invalidate and refetch these endpoints on success. Refetches are awaited
+   * before `onSuccess` runs and before the mutation settles, so `isPending`
+   * (spinners, disabled buttons, open confirm modals) lasts until the UI is
+   * consistent with the mutation — active queries on these endpoints become
+   * part of the mutation's perceived duration.
+   *
+   * On every render, useMutation overwrites the options on any mutation
+   * that is still running, and it doesn't look up onSuccess (which does the
+   * invalidating) until the request resolves. So a mutation invalidates the
+   * list from the most recent render, not the one from when `mutate` was
+   * called. Adding items across renders is harmless (extra invalidations at
+   * worst), but don't remove items and assume they'll still be queued up for
+   * invalidation.
+   */
+  invalidateEndpoints?: readonly (keyof typeof api)[]
+}
+
 export const useApiMutation = <Params, Data>(
   f: (p: Params, fp: FetchParams) => Promise<ApiResult<Data>>,
-  options?: Omit<
-    // __signal bit makes it so you can pass a signal to mutate and mutateAsync.
-    // the underscores make it virtually impossible for this to conflict with an
-    // actual API field
-    UseMutationOptions<Data, ApiError, Params & { __signal?: AbortSignal }>,
-    'mutationFn' | 'onSettled'
-  >
-) =>
-  useMutation({
+  options?: ApiMutationOptions<Params, Data>
+) => {
+  const { invalidateEndpoints, onSuccess, ...mutationOptions } = options ?? {}
+  const onSuccessWithInvalidation = invalidateEndpoints?.length
+    ? async (...args: Parameters<NonNullable<typeof onSuccess>>) => {
+        await Promise.all(invalidateEndpoints.map((e) => queryClient.invalidateEndpoint(e)))
+        await onSuccess?.(...args)
+      }
+    : onSuccess
+
+  return useMutation({
     mutationFn: ({ __signal, ...params }) =>
+      // __signal bit makes it so you can pass a signal to mutate and mutateAsync.
+      // the underscores make it virtually impossible for this to conflict with an
+      // actual API field.
       // Pretty safe cast: signal is an optional addition at the call site, not
       // part of the original Params type. Removing it via destructuring gives
       // us back Params, but TS can't prove Omit<Params & {signal?}, 'signal'>
@@ -400,5 +427,7 @@ export const useApiMutation = <Params, Data>(
           throw result.data
         }),
     // no catch, let unexpected errors bubble up
-    ...options,
+    ...mutationOptions,
+    onSuccess: onSuccessWithInvalidation,
   })
+}

app/forms/floating-ip-create.tsx

@@ -60,9 +60,8 @@ export default function CreateFloatingIpSideModalForm() {
   const navigate = useNavigate()
 
   const createFloatingIp = useApiMutation(api.floatingIpCreate, {
+    invalidateEndpoints: ['floatingIpList', 'systemIpPoolUtilizationView'],
     onSuccess(floatingIp) {
-      queryClient.invalidateEndpoint('floatingIpList')
-      queryClient.invalidateEndpoint('systemIpPoolUtilizationView')
       // prettier-ignore
       addToast(<>Floating IP <HL>{floatingIp.name}</HL> created</>)
       navigate(pb.floatingIps(projectSelector))

app/pages/system/silos/SiloIpPoolsTab.tsx

@@ -162,15 +162,11 @@ export default function SiloIpPoolsTab() {
   )
 
   const { mutateAsync: updatePoolLink } = useApiMutation(api.systemIpPoolSiloUpdate, {
-    onSuccess() {
-      queryClient.invalidateEndpoint('siloIpPoolList')
-      queryClient.invalidateEndpoint('systemIpPoolSiloList')
-    },
+    invalidateEndpoints: ['siloIpPoolList', 'systemIpPoolSiloList'],
   })
   const { mutateAsync: unlinkPool } = useApiMutation(api.systemIpPoolSiloUnlink, {
+    invalidateEndpoints: ['siloIpPoolList', 'systemIpPoolSiloList'],
     onSuccess() {
-      queryClient.invalidateEndpoint('siloIpPoolList')
-      queryClient.invalidateEndpoint('systemIpPoolSiloList')
       // We only have the ID, so will show a generic confirmation message
       addToast({ content: 'IP pool unlinked' })
     },