From 7fa0832edd5b814878991875dc3a15d4b2ef3427 Mon Sep 17 00:00:00 2001 From: Andrew Chang <69671930+Andrushika@users.noreply.github.com> Date: Tue, 4 Aug 2026 03:14:55 +0800 Subject: [PATCH] [v3-3-test] UI: Refresh task details immediately when switching tasks (#70789) Cached task data can remain fresh for five minutes after navigating away, leaving the selected task out of sync with Graph view until the next polling interval. (cherry picked from commit a14105b25926246a8d61ecb494c880b477e6f862) Co-authored-by: Andrew Chang <69671930+Andrushika@users.noreply.github.com> --- .../ui/src/components/TaskTrySelect.test.tsx | 128 ++++++++++++++++ .../ui/src/components/TaskTrySelect.tsx | 1 + .../pages/TaskInstance/TaskInstance.test.tsx | 144 ++++++++++++++++++ .../src/pages/TaskInstance/TaskInstance.tsx | 1 + 4 files changed, 274 insertions(+) create mode 100644 airflow-core/src/airflow/ui/src/components/TaskTrySelect.test.tsx create mode 100644 airflow-core/src/airflow/ui/src/pages/TaskInstance/TaskInstance.test.tsx 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, }, );