From abccac89dc29d9a02a2184d11af1846578571e19 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Mon, 23 Mar 2026 22:37:24 -0700
Subject: [PATCH] Address greptile review feedback (greploop iteration 3)

- Add tests for getTestSourceOptions (2 tests)
- Add null guards to selectAntOption test helper

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../policies/pipeline_test_drawer.test.tsx    |  6 +++--
 .../policies/pipeline_utils.test.ts           | 23 +++++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.test.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.test.tsx
index f4bb3837fdb..9ebee612087 100644
--- a/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.test.tsx
+++ b/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.test.tsx
@@ -95,12 +95,14 @@ const defaultProps = {
 };
 
 async function selectAntOption(user: ReturnType<typeof userEvent.setup>, title: string) {
-  const selector = document.querySelector(".ant-select-selector") as HTMLElement;
+  const selector = document.querySelector(".ant-select-selector");
+  if (!selector) throw new Error("Ant Select selector not found");
   await user.click(selector);
   await waitFor(() => {
     expect(document.querySelector(`[title="${title}"].ant-select-item-option`)).toBeInTheDocument();
   });
-  const option = document.querySelector(`[title="${title}"].ant-select-item-option`) as HTMLElement;
+  const option = document.querySelector(`[title="${title}"].ant-select-item-option`);
+  if (!option) throw new Error(`Ant Select option "${title}" not found`);
   await user.click(option);
 }
 
diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_utils.test.ts b/ui/litellm-dashboard/src/components/policies/pipeline_utils.test.ts
index 1e0c191fd5a..59ca24dca48 100644
--- a/ui/litellm-dashboard/src/components/policies/pipeline_utils.test.ts
+++ b/ui/litellm-dashboard/src/components/policies/pipeline_utils.test.ts
@@ -7,6 +7,7 @@ import {
   derivePipelineFromPolicy,
   complianceMatchExpected,
   getPromptsForTestSource,
+  getTestSourceOptions,
   TEST_SOURCE_QUICK,
   TEST_SOURCE_ALL,
 } from "./pipeline_utils";
@@ -219,3 +220,25 @@ describe("getPromptsForTestSource", () => {
     vi.restoreAllMocks();
   });
 });
+
+describe("getTestSourceOptions", () => {
+  it("should always include quick chat and all datasets options", () => {
+    vi.spyOn(complianceData, "getFrameworks").mockReturnValue([]);
+    const options = getTestSourceOptions();
+    expect(options[0]).toEqual({ value: TEST_SOURCE_QUICK, label: "Quick chat (custom message)" });
+    expect(options[options.length - 1]).toEqual({ value: TEST_SOURCE_ALL, label: "All compliance datasets" });
+    vi.restoreAllMocks();
+  });
+
+  it("should include framework names from getFrameworks", () => {
+    vi.spyOn(complianceData, "getFrameworks").mockReturnValue([
+      { name: "GDPR", icon: "", description: "", categories: [] },
+      { name: "EU AI Act", icon: "", description: "", categories: [] },
+    ]);
+    const options = getTestSourceOptions();
+    expect(options).toHaveLength(4);
+    expect(options[1]).toEqual({ value: "GDPR", label: "GDPR" });
+    expect(options[2]).toEqual({ value: "EU AI Act", label: "EU AI Act" });
+    vi.restoreAllMocks();
+  });
+});