From 50e24512e04840caf1345ada0d6f0308fd91fdbf Mon Sep 17 00:00:00 2001 From: Shivam Date: Wed, 29 Jul 2026 20:21:25 -0700 Subject: [PATCH] UI: Keep task try history consistent when switching tasks Task try history responses can include duplicate or non-attempt records, which caused the UI to display mixed or invalid attempts after task navigation. --- .../ui/src/components/TaskTrySelect.test.tsx | 206 +++++++++++++++--- .../ui/src/components/TaskTrySelect.tsx | 12 +- 2 files changed, 180 insertions(+), 38 deletions(-) diff --git a/airflow-core/src/airflow/ui/src/components/TaskTrySelect.test.tsx b/airflow-core/src/airflow/ui/src/components/TaskTrySelect.test.tsx index 9a1e39766b11a..142dc5fc36c7d 100644 --- a/airflow-core/src/airflow/ui/src/components/TaskTrySelect.test.tsx +++ b/airflow-core/src/airflow/ui/src/components/TaskTrySelect.test.tsx @@ -18,8 +18,8 @@ */ 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 { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import type { PropsWithChildren, ReactNode } from "react"; import { MemoryRouter } from "react-router-dom"; import { afterEach, describe, expect, it, vi } from "vitest"; @@ -33,6 +33,12 @@ import { import { TaskTrySelect } from "./TaskTrySelect"; +vi.mock("src/components/StateBadge", () => ({ + StateBadge: ({ children, state }: { readonly children?: ReactNode; readonly state?: string | null }) => ( + {children} + ), +})); + vi.mock("src/utils", async () => { const actual = await vi.importActual("src/utils"); @@ -47,34 +53,59 @@ const DAG_RUN_ID = "test_run"; const TASK_A = "task_a"; const TASK_B = "task_b"; -const buildTaskInstance = (taskId: string, tryNumber: number): TaskInstanceResponse => +const buildTaskInstance = ( + taskId: string, + tryNumber: number, + { + mapIndex = -1, + state = "success", + }: { + readonly mapIndex?: number; + readonly state?: TaskInstanceResponse["state"]; + } = {}, +): TaskInstanceResponse => ({ dag_id: DAG_ID, dag_run_id: DAG_RUN_ID, - id: `${taskId}-id`, - map_index: -1, - state: "success", + id: `${taskId}-${mapIndex}`, + map_index: mapIndex, + state, 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 buildTaskTry = ( + taskId: string, + tryNumber: number, + { + mapIndex = -1, + state = "success", + }: { + readonly mapIndex?: number; + readonly state?: TaskInstanceHistoryResponse["state"]; + } = {}, +): TaskInstanceHistoryResponse => ({ + ...buildTaskInstance(taskId, tryNumber, { mapIndex, state }), +}); -const buildTaskTries = (tryNumbers: Array): TaskInstanceHistoryCollectionResponse => ({ - task_instances: tryNumbers.map(buildTaskTry), - total_entries: tryNumbers.length, +const buildTaskTries = ( + taskInstances: Array, +): TaskInstanceHistoryCollectionResponse => ({ + task_instances: taskInstances, + total_entries: taskInstances.length, }); +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + staleTime: 5 * 60 * 1000, + }, + }, + }); + const createWrapper = (queryClient: QueryClient) => ({ children }: PropsWithChildren) => ( @@ -85,18 +116,31 @@ const createWrapper = ); +const expectTries = async (tries: Array) => { + await waitFor(() => { + expect( + screen + .getAllByTestId(/^log-attempt-select-button-/u) + .map((button) => button.getAttribute("data-testid")), + ).toEqual(tries.map((tryNumber) => `log-attempt-select-button-${tryNumber}`)); + }); + expect(screen.queryByTestId("log-attempt-select-button-0")).toBeNull(); +}; + +const expectTryState = (tryNumber: number, state: string) => { + expect( + screen + .getByTestId(`log-attempt-select-button-${tryNumber}`) + .querySelector("[data-state]") + ?.getAttribute("data-state"), + ).toBe(state); +}; + 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 queryClient = createQueryClient(); const params = { dagId: DAG_ID, dagRunId: DAG_RUN_ID, @@ -106,9 +150,11 @@ describe("TaskTrySelect", () => { queryClient.setQueryData( UseTaskInstanceServiceGetMappedTaskInstanceTriesKeyFn(params), - buildTaskTries([1, 2]), + buildTaskTries([1, 2].map((tryNumber) => buildTaskTry(TASK_A, tryNumber))), + ); + vi.spyOn(TaskInstanceService, "getMappedTaskInstanceTries").mockResolvedValue( + buildTaskTries([1, 2, 3].map((tryNumber) => buildTaskTry(TASK_A, tryNumber))), ); - vi.spyOn(TaskInstanceService, "getMappedTaskInstanceTries").mockResolvedValue(buildTaskTries([1, 2, 3])); const { rerender } = render( , @@ -117,12 +163,102 @@ describe("TaskTrySelect", () => { 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"]); + await expectTries([1, 2, 3]); expect(TaskInstanceService.getMappedTaskInstanceTries).toHaveBeenCalledWith(params); }); + + it("keeps positive tries unique while switching between task instances", async () => { + const queryClient = createQueryClient(); + const histories = { + start: buildTaskTries([1, 2, 2, 3, 4].map((tryNumber) => buildTaskTry("start", tryNumber))), + task_1: buildTaskTries([ + buildTaskTry("task_1", 0, { state: "skipped" }), + buildTaskTry("task_1", 1), + buildTaskTry("task_1", 2), + ]), + task_2: buildTaskTries([buildTaskTry("task_2", 1), buildTaskTry("task_2", 2)]), + }; + + const getTries = vi + .spyOn(TaskInstanceService, "getMappedTaskInstanceTries") + .mockResolvedValue(histories.start); + + const { rerender } = render(, { + wrapper: createWrapper(queryClient), + }); + + await expectTries([1, 2, 3, 4]); + + getTries.mockResolvedValue(histories.task_1); + rerender(); + await expectTries([1, 2]); + + getTries.mockResolvedValue(histories.task_2); + rerender(); + await expectTries([1, 2]); + + getTries.mockResolvedValue(histories.task_1); + rerender(); + await expectTries([1, 2]); + + getTries.mockResolvedValue(histories.start); + rerender(); + await expectTries([1, 2, 3, 4]); + }); + + it("uses a real current try but not retry or null placeholders", async () => { + const queryClient = createQueryClient(); + const params = { + dagId: DAG_ID, + dagRunId: DAG_RUN_ID, + mapIndex: 1, + taskId: "mapped_task", + }; + const history = buildTaskTries([ + buildTaskTry("mapped_task", 1, { mapIndex: 1 }), + buildTaskTry("mapped_task", 2, { mapIndex: 1, state: "failed" }), + ]); + const onSelectTryNumber = vi.fn(); + + queryClient.setQueryData(UseTaskInstanceServiceGetMappedTaskInstanceTriesKeyFn(params), history); + vi.spyOn(TaskInstanceService, "getMappedTaskInstanceTries").mockResolvedValue(history); + + const { rerender } = render( + , + { wrapper: createWrapper(queryClient) }, + ); + + await expectTries([1, 2]); + expectTryState(2, "success"); + + rerender( + , + ); + expectTryState(2, "failed"); + + rerender( + , + ); + expectTryState(2, "failed"); + + fireEvent.click(screen.getByTestId("log-attempt-select-button-1")); + expect(onSelectTryNumber).toHaveBeenCalledOnce(); + expect(onSelectTryNumber).toHaveBeenCalledWith(1); + expect(TaskInstanceService.getMappedTaskInstanceTries).toHaveBeenCalledWith(params); + + rerender(); + expect(screen.queryByTestId(/^log-attempt-select-button-/u)).toBeNull(); + }); }); diff --git a/airflow-core/src/airflow/ui/src/components/TaskTrySelect.tsx b/airflow-core/src/airflow/ui/src/components/TaskTrySelect.tsx index 8cbf595d0c288..75132ffa013b3 100644 --- a/airflow-core/src/airflow/ui/src/components/TaskTrySelect.tsx +++ b/airflow-core/src/airflow/ui/src/components/TaskTrySelect.tsx @@ -71,11 +71,17 @@ export const TaskTrySelect = ({ onSelectTryNumber, selectedTryNumber, taskInstan const logAttemptDropdownLimit = 10; const showDropdown = finalTryNumber > logAttemptDropdownLimit; - // For some reason tries aren't sorted by try_number - const sortedTries = [...(tiHistory?.task_instances ?? [])].sort( - (tryA, tryB) => tryA.try_number - tryB.try_number, + const triesByNumber = new Map( + (tiHistory?.task_instances ?? []).filter((ti) => ti.try_number > 0).map((ti) => [ti.try_number, ti]), ); + if (finalTryNumber > 0 && state !== "up_for_retry" && state !== null) { + // The current task instance is authoritative when it is also present in history. + triesByNumber.set(finalTryNumber, taskInstance); + } + + const sortedTries = [...triesByNumber.values()].sort((tryA, tryB) => tryA.try_number - tryB.try_number); + const tryOptions = createListCollection({ items: sortedTries.map((ti) => ({ task_instance: ti,