diff --git a/airflow-core/src/airflow/ui/src/components/TaskTrySelect.test.tsx b/airflow-core/src/airflow/ui/src/components/TaskTrySelect.test.tsx new file mode 100644 index 0000000000000..9a1e39766b11a --- /dev/null +++ b/airflow-core/src/airflow/ui/src/components/TaskTrySelect.test.tsx @@ -0,0 +1,128 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { ChakraProvider, defaultSystem } from "@chakra-ui/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen } from "@testing-library/react"; +import type { PropsWithChildren } from "react"; +import { MemoryRouter } from "react-router-dom"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +import { UseTaskInstanceServiceGetMappedTaskInstanceTriesKeyFn } from "openapi/queries"; +import { + TaskInstanceService, + type TaskInstanceHistoryCollectionResponse, + type TaskInstanceHistoryResponse, + type TaskInstanceResponse, +} from "openapi/requests"; + +import { TaskTrySelect } from "./TaskTrySelect"; + +vi.mock("src/utils", async () => { + const actual = await vi.importActual("src/utils"); + + return { + ...actual, + useAutoRefresh: vi.fn(() => false), + }; +}); + +const DAG_ID = "test_dag"; +const DAG_RUN_ID = "test_run"; +const TASK_A = "task_a"; +const TASK_B = "task_b"; + +const buildTaskInstance = (taskId: string, tryNumber: number): TaskInstanceResponse => + ({ + dag_id: DAG_ID, + dag_run_id: DAG_RUN_ID, + id: `${taskId}-id`, + map_index: -1, + state: "success", + task_display_name: taskId, + task_id: taskId, + try_number: tryNumber, + }) as TaskInstanceResponse; + +const buildTaskTry = (tryNumber: number): TaskInstanceHistoryResponse => + ({ + dag_id: DAG_ID, + dag_run_id: DAG_RUN_ID, + map_index: -1, + state: "success", + task_display_name: TASK_A, + task_id: TASK_A, + try_number: tryNumber, + }) as TaskInstanceHistoryResponse; + +const buildTaskTries = (tryNumbers: Array): TaskInstanceHistoryCollectionResponse => ({ + task_instances: tryNumbers.map(buildTaskTry), + total_entries: tryNumbers.length, +}); + +const createWrapper = + (queryClient: QueryClient) => + ({ children }: PropsWithChildren) => ( + + + {children} + + + ); + +afterEach(() => vi.restoreAllMocks()); + +describe("TaskTrySelect", () => { + it("refetches cached tries immediately when switching tasks", async () => { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + staleTime: 5 * 60 * 1000, + }, + }, + }); + const params = { + dagId: DAG_ID, + dagRunId: DAG_RUN_ID, + mapIndex: -1, + taskId: TASK_A, + }; + + queryClient.setQueryData( + UseTaskInstanceServiceGetMappedTaskInstanceTriesKeyFn(params), + buildTaskTries([1, 2]), + ); + vi.spyOn(TaskInstanceService, "getMappedTaskInstanceTries").mockResolvedValue(buildTaskTries([1, 2, 3])); + + const { rerender } = render( + , + { wrapper: createWrapper(queryClient) }, + ); + + rerender(); + + expect(await screen.findByTestId("log-attempt-select-button-3")).toBeTruthy(); + expect( + screen + .getAllByTestId(/^log-attempt-select-button-/u) + .map((button) => button.getAttribute("data-testid")), + ).toEqual(["log-attempt-select-button-1", "log-attempt-select-button-2", "log-attempt-select-button-3"]); + expect(TaskInstanceService.getMappedTaskInstanceTries).toHaveBeenCalledWith(params); + }); +}); diff --git a/airflow-core/src/airflow/ui/src/components/TaskTrySelect.tsx b/airflow-core/src/airflow/ui/src/components/TaskTrySelect.tsx index 5115383820ebb..8cbf595d0c288 100644 --- a/airflow-core/src/airflow/ui/src/components/TaskTrySelect.tsx +++ b/airflow-core/src/airflow/ui/src/components/TaskTrySelect.tsx @@ -60,6 +60,7 @@ export const TaskTrySelect = ({ onSelectTryNumber, selectedTryNumber, taskInstan query.state.data?.task_instances.some((ti) => isStatePending(ti.state)) || isStatePending(state) ? refetchInterval : false, + staleTime: 0, }, ); diff --git a/airflow-core/src/airflow/ui/src/pages/TaskInstance/TaskInstance.test.tsx b/airflow-core/src/airflow/ui/src/pages/TaskInstance/TaskInstance.test.tsx new file mode 100644 index 0000000000000..7c56371f09825 --- /dev/null +++ b/airflow-core/src/airflow/ui/src/pages/TaskInstance/TaskInstance.test.tsx @@ -0,0 +1,144 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { fireEvent, render, screen } from "@testing-library/react"; +import type { PropsWithChildren } from "react"; +import { Link, MemoryRouter, Route, Routes } from "react-router-dom"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +import { UseTaskInstanceServiceGetMappedTaskInstanceKeyFn } from "openapi/queries"; +import { TaskInstanceService, type TaskInstanceResponse } from "openapi/requests"; + +import { TaskInstance } from "./TaskInstance"; + +vi.mock("src/hooks/useHITLReviewTabs", () => ({ + useHITLReviewTabs: vi.fn(() => ({ tabs: [] })), +})); +vi.mock("src/hooks/usePluginTabs", () => ({ + usePluginTabs: vi.fn(() => []), +})); +vi.mock("src/hooks/useRequiredActionTabs", () => ({ + useRequiredActionTabs: vi.fn(() => ({ tabs: [] })), +})); +vi.mock("src/layouts/Details/DetailsLayout", () => ({ + DetailsLayout: ({ children }: PropsWithChildren) => children, +})); +vi.mock("src/queries/useGridTISummaries.ts", () => ({ + useGridTiSummariesStream: vi.fn(() => ({ summariesByRunId: new Map() })), +})); +vi.mock("src/utils", async () => { + const actual = await vi.importActual("src/utils"); + + return { + ...actual, + useAutoRefresh: vi.fn(() => false), + useDocumentTitle: vi.fn(), + }; +}); +vi.mock("./Header", () => ({ + Header: ({ taskInstance }: { readonly taskInstance: TaskInstanceResponse }) => ( +
+ {taskInstance.task_id}:{taskInstance.state ?? "none"}:{taskInstance.try_number} +
+ ), +})); + +const DAG_ID = "test_dag"; +const DAG_RUN_ID = "test_run"; +const TASK_A = "task_a"; +const TASK_B = "task_b"; + +const buildTaskInstance = ( + taskId: string, + state: TaskInstanceResponse["state"], + tryNumber: number, +): TaskInstanceResponse => + ({ + dag_id: DAG_ID, + dag_run_id: DAG_RUN_ID, + id: `${taskId}-id`, + map_index: -1, + state, + task_display_name: taskId, + task_id: taskId, + try_number: tryNumber, + }) as TaskInstanceResponse; + +const buildTaskInstanceKey = (taskId: string) => + UseTaskInstanceServiceGetMappedTaskInstanceKeyFn({ + dagId: DAG_ID, + dagRunId: DAG_RUN_ID, + mapIndex: -1, + taskId, + }); + +const createWrapper = + (queryClient: QueryClient) => + ({ children }: PropsWithChildren) => ( + {children} + ); + +afterEach(() => vi.restoreAllMocks()); + +describe("TaskInstance", () => { + it("refetches a cached task instance immediately when switching tasks", async () => { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + staleTime: 5 * 60 * 1000, + }, + }, + }); + const cachedTaskA = buildTaskInstance(TASK_A, null, 2); + const latestTaskA = buildTaskInstance(TASK_A, "success", 3); + const taskB = buildTaskInstance(TASK_B, "success", 1); + + queryClient.setQueryData(buildTaskInstanceKey(TASK_A), cachedTaskA); + queryClient.setQueryData(buildTaskInstanceKey(TASK_B), taskB); + vi.spyOn(TaskInstanceService, "getMappedTaskInstance").mockImplementation( + ({ taskId }) => + Promise.resolve(taskId === TASK_A ? latestTaskA : taskB) as unknown as ReturnType< + typeof TaskInstanceService.getMappedTaskInstance + >, + ); + + render( + + Open task A + + } path="/dags/:dagId/runs/:runId/tasks/:taskId" /> + + , + { wrapper: createWrapper(queryClient) }, + ); + + expect(await screen.findByText(`${TASK_B}:success:1`)).toBeTruthy(); + + fireEvent.click(screen.getByRole("link", { name: "Open task A" })); + + expect(await screen.findByText(`${TASK_A}:success:3`)).toBeTruthy(); + expect(TaskInstanceService.getMappedTaskInstance).toHaveBeenCalledWith({ + dagId: DAG_ID, + dagRunId: DAG_RUN_ID, + mapIndex: -1, + taskId: TASK_A, + }); + }); +}); diff --git a/airflow-core/src/airflow/ui/src/pages/TaskInstance/TaskInstance.tsx b/airflow-core/src/airflow/ui/src/pages/TaskInstance/TaskInstance.tsx index 80fdd93268e34..75371ce62707b 100644 --- a/airflow-core/src/airflow/ui/src/pages/TaskInstance/TaskInstance.tsx +++ b/airflow-core/src/airflow/ui/src/pages/TaskInstance/TaskInstance.tsx @@ -82,6 +82,7 @@ export const TaskInstance = () => { { enabled: !isNaN(parsedMapIndex), refetchInterval: (query) => (isStatePending(query.state.data?.state) ? refetchInterval : false), + staleTime: 0, }, );