From eb831d956ccb328411ebb86a0161ac7a23b4aba8 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 17 Sep 2026 23:46:53 -0700
Subject: [PATCH] test(ui): address review feedback

---
 tests/e2e/ui/helpers/roundTrip.ts             | 23 +++++++++---
 tests/e2e/ui/tests/prompts/addPrompt.spec.ts  |  4 +--
 .../tests/tagManagement/tagManagement.spec.ts |  9 ++---
 ...PaginatedSearchSelect.integration.test.tsx | 36 -------------------
 4 files changed, 26 insertions(+), 46 deletions(-)

diff --git a/tests/e2e/ui/helpers/roundTrip.ts b/tests/e2e/ui/helpers/roundTrip.ts
index 4175ba6ec72..1fc2d0aec1a 100644
--- a/tests/e2e/ui/helpers/roundTrip.ts
+++ b/tests/e2e/ui/helpers/roundTrip.ts
@@ -46,11 +46,26 @@ export async function runWithCleanup(
   try {
     if (outcome.status === "failure") throw outcome.error;
   } finally {
-    const cleanupSucceeded = await Promise.resolve()
+    const cleanupOutcome = await Promise.resolve()
       .then(cleanup)
-      .catch(() => false);
-    if (outcome.status === "success" && !cleanupSucceeded) {
-      throw new Error("Failed to clean up UI E2E resource");
+      .then(
+        (succeeded) =>
+          succeeded
+            ? { status: "success" as const }
+            : {
+                status: "failure" as const,
+                error: new Error("Failed to clean up UI E2E resource"),
+              },
+        (error: unknown) => ({ status: "failure" as const, error }),
+      );
+    if (cleanupOutcome.status === "failure") {
+      if (outcome.status === "failure") {
+        throw new AggregateError(
+          [outcome.error, cleanupOutcome.error],
+          "Action and cleanup failed",
+        );
+      }
+      throw cleanupOutcome.error;
     }
   }
 }
diff --git a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts
index 891739fda28..9d85236c4a6 100644
--- a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts
+++ b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts
@@ -27,7 +27,7 @@ test.describe("Prompt upload form", () => {
           name: "e2e.prompt",
           mimeType: "text/plain",
           buffer: Buffer.from(
-            `model: fake-openai-gpt-4\ntemplate: "${promptContent}"\n`,
+            `---\nmodel: fake-openai-gpt-4\n---\n${promptContent}\n`,
           ),
         });
         await expect(page.getByText("Selected: e2e.prompt")).toBeVisible();
@@ -58,7 +58,7 @@ test.describe("Prompt upload form", () => {
             };
             return promptInfo.raw_prompt_template?.content;
           })
-          .toContain(promptContent);
+          .toBe(promptContent);
         await expect(page.getByText(promptId, { exact: true })).toBeVisible();
       },
       async () => {
diff --git a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts
index bf46b5ea161..fe659080eab 100644
--- a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts
+++ b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts
@@ -31,10 +31,11 @@ test.describe("Tag management", () => {
 
         await expect
           .poll(async () => {
-            const response = await readBack<
-              Record<string, Record<string, unknown>>
-            >(page, "/tag/list");
-            return Object.values(response).some((tag) => tag.name === tagName);
+            const response = await readBack<Array<{ name: string }>>(
+              page,
+              "/tag/list",
+            );
+            return response.some((tag) => tag.name === tagName);
           })
           .toBe(true);
         await expect(page.getByText(tagName, { exact: true })).toBeVisible();
diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx
index 905610a77e6..2b948ca8420 100644
--- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx
+++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.integration.test.tsx
@@ -1,6 +1,5 @@
 import { fireEvent, renderWithProviders as render, screen, waitFor } from "../../../tests/test-utils";
 import userEvent from "@testing-library/user-event";
-import { useQuery } from "@tanstack/react-query";
 import { useState } from "react";
 import { describe, expect, it, vi } from "vitest";
 
@@ -407,39 +406,4 @@ describe("PaginatedSearchSelect", () => {
     expect(input).toHaveValue("aliasalpha");
     await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith("aliasalpha"));
   });
-
-  it("keeps the latest query results when an earlier response resolves last", async () => {
-    const pending = new Map<string, (options: SearchSelectOption[]) => void>();
-
-    function QueryBackedSelect() {
-      const [query, setQuery] = useState("");
-      const result = useQuery({
-        queryKey: ["paginated-select-race", query],
-        queryFn: () =>
-          new Promise<SearchSelectOption[]>((resolve) => {
-            pending.set(query, resolve);
-          }),
-        enabled: query.length > 0,
-      });
-      return <PaginatedSearchSelect options={result.data ?? []} onValueChange={vi.fn()} onSearchChange={setQuery} />;
-    }
-
-    const user = userEvent.setup();
-    render(<QueryBackedSelect />);
-    const input = screen.getByRole("combobox");
-    await user.click(input);
-    await user.type(input, "A");
-    await waitFor(() => expect(pending.has("A")).toBe(true));
-    await user.clear(input);
-    await user.type(input, "B");
-    await waitFor(() => expect(pending.has("B")).toBe(true));
-
-    pending.get("B")?.([{ label: "B result", value: "b" }]);
-    await user.click(screen.getByRole("combobox"));
-    expect(await screen.findByText("B result")).toBeInTheDocument();
-
-    pending.get("A")?.([{ label: "A result", value: "a" }]);
-    await waitFor(() => expect(screen.queryByText("A result")).not.toBeInTheDocument());
-    expect(screen.getByText("B result")).toBeInTheDocument();
-  });
 });