diff --git a/ui-v2/src/api/flow-runs/flow-runs.test.ts b/ui-v2/src/api/flow-runs/flow-runs.test.ts index 77e5791a4932..f8eba299f23a 100644 --- a/ui-v2/src/api/flow-runs/flow-runs.test.ts +++ b/ui-v2/src/api/flow-runs/flow-runs.test.ts @@ -1,11 +1,15 @@ import type { components } from "@/api/prefect"; import { createFakeFlowRun } from "@/mocks"; import { QueryClient, useSuspenseQuery } from "@tanstack/react-query"; -import { renderHook, waitFor } from "@testing-library/react"; +import { act, renderHook, waitFor } from "@testing-library/react"; import { buildApiUrl, createWrapper, server } from "@tests/utils"; import { http, HttpResponse } from "msw"; import { describe, expect, it } from "vitest"; -import { buildListFlowRunsQuery } from "."; +import { + buildListFlowRunsQuery, + queryKeyFactory, + useDeploymentCreateFlowRun, +} from "."; type FlowRun = components["schemas"]["FlowRun"]; @@ -17,6 +21,13 @@ describe("flow runs api", () => { }), ); }; + const mockCreateDeploymentFlowRunAPI = (flowRun: FlowRun) => { + server.use( + http.post(buildApiUrl("/deployments/:id/create_flow_run"), () => { + return HttpResponse.json(flowRun); + }), + ); + }; describe("flowRunsQueryParams", () => { it("fetches paginated flow runs with default parameters", async () => { @@ -72,4 +83,54 @@ describe("flow runs api", () => { expect(refetchInterval).toBe(customRefetchInterval); }); }); + describe("useDeploymentCreateFlowRun", () => { + it("invalidates cache and fetches updated value", async () => { + const FILTER = { + sort: "ID_DESC", + offset: 0, + } as const; + const queryClient = new QueryClient(); + const EXISTING_CACHE = [createFakeFlowRun(), createFakeFlowRun()]; + const MOCK_NEW_DATA_ID = "2"; + const NEW_FLOW_RUN_DATA = createFakeFlowRun({ id: MOCK_NEW_DATA_ID }); + + // ------------ Mock API requests after queries are invalidated + const mockData = [NEW_FLOW_RUN_DATA, ...EXISTING_CACHE]; + mockCreateDeploymentFlowRunAPI(NEW_FLOW_RUN_DATA); + mockFetchFlowRunsAPI(mockData); + + // ------------ Initialize cache + queryClient.setQueryData(queryKeyFactory.list(FILTER), EXISTING_CACHE); + + // ------------ Initialize hooks to test + const { result: useDeploymentCreateFlowRunResult } = renderHook( + useDeploymentCreateFlowRun, + { wrapper: createWrapper({ queryClient }) }, + ); + + const { result: useListFlowRunsResult } = renderHook( + () => useSuspenseQuery(buildListFlowRunsQuery(FILTER)), + { wrapper: createWrapper({ queryClient }) }, + ); + + // ------------ Invoke mutation + act(() => + useDeploymentCreateFlowRunResult.current.createDeploymentFlowRun({ + id: MOCK_NEW_DATA_ID, + }), + ); + + // ------------ Assert + await waitFor(() => + expect(useDeploymentCreateFlowRunResult.current.isSuccess).toBe(true), + ); + + expect(useListFlowRunsResult.current.data).toHaveLength(3); + + const newFlowRun = useListFlowRunsResult.current.data?.find( + (flowRun) => flowRun.id === MOCK_NEW_DATA_ID, + ); + expect(newFlowRun).toMatchObject(NEW_FLOW_RUN_DATA); + }); + }); }); diff --git a/ui-v2/src/api/flow-runs/index.ts b/ui-v2/src/api/flow-runs/index.ts index 6b55ee7f0da2..5c153afe09e7 100644 --- a/ui-v2/src/api/flow-runs/index.ts +++ b/ui-v2/src/api/flow-runs/index.ts @@ -1,4 +1,8 @@ -import { queryOptions } from "@tanstack/react-query"; +import { + queryOptions, + useMutation, + useQueryClient, +} from "@tanstack/react-query"; import { Deployment } from "../deployments"; import { Flow } from "../flows"; import { components } from "../prefect"; @@ -68,3 +72,60 @@ export const buildListFlowRunsQuery = ( refetchInterval, }); }; + +// ----- ✍🏼 Mutations 🗄️ +// ---------------------------- + +type MutateCreateFlowRun = { + id: string; +} & components["schemas"]["DeploymentFlowRunCreate"]; +/** + * Hook for creating a new flow run from an automation + * + * @returns Mutation object for creating a flow run with loading/error states and trigger function + * + * @example + * ```ts + * const { createDeploymentFlowRun, isLoading } = useDeploymentCreateFlowRun(); + * + * createDeploymentFlowRun({ deploymentId, ...body }, { + * onSuccess: () => { + * // Handle successful creation + * console.log('Flow run created successfully'); + * }, + * onError: (error) => { + * // Handle error + * console.error('Failed to create flow run:', error); + * } + * }); + * ``` + */ +export const useDeploymentCreateFlowRun = () => { + const queryClient = useQueryClient(); + const { mutate: createDeploymentFlowRun, ...rest } = useMutation({ + mutationFn: async ({ id, ...body }: MutateCreateFlowRun) => { + const res = await getQueryService().POST( + "/deployments/{id}/create_flow_run", + { + body, + params: { path: { id } }, + }, + ); + + if (!res.data) { + throw new Error("'data' expected"); + } + return res.data; + }, + onSuccess: () => { + // After a successful creation, invalidate only list queries to refetch + return queryClient.invalidateQueries({ + queryKey: queryKeyFactory.lists(), + }); + }, + }); + return { + createDeploymentFlowRun, + ...rest, + }; +};