From 3daf6386bc053db72f57a40c8e8d9b6dda2414a5 Mon Sep 17 00:00:00 2001 From: Andrew Chang Date: Fri, 31 Jul 2026 02:20:54 +0800 Subject: [PATCH] UI: Refresh task details immediately when switching tasks 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. --- .../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, }, );