From 45a54d64cd86f36d480304ab6cba6c79dce96551 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Mon, 23 Mar 2026 21:58:25 -0700
Subject: [PATCH] Address review feedback: rename, fix FailIcon, add missing
 tests
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Rename pipeline_icons.tsx to pipeline_utils.tsx (better reflects mixed
  contents: constants, helpers, and icons)
- Fix FailIcon missing cross mark paths (pre-existing bug, now renders ✕)
- Add 4 tests for getPromptsForTestSource covering all branches

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../policies/pipeline_flow_builder.tsx        |  2 +-
 .../policies/pipeline_test_drawer.tsx         |  2 +-
 ...e_icons.test.ts => pipeline_utils.test.ts} | 40 ++++++++++++++++++-
 ...{pipeline_icons.tsx => pipeline_utils.tsx} |  2 +
 4 files changed, 42 insertions(+), 4 deletions(-)
 rename ui/litellm-dashboard/src/components/policies/{pipeline_icons.test.ts => pipeline_utils.test.ts} (77%)
 rename ui/litellm-dashboard/src/components/policies/{pipeline_icons.tsx => pipeline_utils.tsx} (99%)

diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx
index feef045b80c..78085052655 100644
--- a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx
+++ b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx
@@ -21,7 +21,7 @@ import {
   PassIcon,
   FailIcon,
   Connector,
-} from "./pipeline_icons";
+} from "./pipeline_utils";
 import { PipelineTestPanel } from "./pipeline_test_drawer";
 
 const { Text } = Typography;
diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.tsx
index d7571d1c6e9..670ff55d480 100644
--- a/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.tsx
+++ b/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.tsx
@@ -11,7 +11,7 @@ import {
   ACTION_LABELS,
   getPromptsForTestSource,
   complianceMatchExpected,
-} from "./pipeline_icons";
+} from "./pipeline_utils";
 
 // ─────────────────────────────────────────────────────────────────────────────
 // Style maps
diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_icons.test.ts b/ui/litellm-dashboard/src/components/policies/pipeline_utils.test.ts
similarity index 77%
rename from ui/litellm-dashboard/src/components/policies/pipeline_icons.test.ts
rename to ui/litellm-dashboard/src/components/policies/pipeline_utils.test.ts
index 3c04bd2fa45..1e0c191fd5a 100644
--- a/ui/litellm-dashboard/src/components/policies/pipeline_icons.test.ts
+++ b/ui/litellm-dashboard/src/components/policies/pipeline_utils.test.ts
@@ -1,4 +1,4 @@
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import {
   createDefaultStep,
   insertStep,
@@ -6,8 +6,12 @@ import {
   updateStepAtIndex,
   derivePipelineFromPolicy,
   complianceMatchExpected,
-} from "./pipeline_icons";
+  getPromptsForTestSource,
+  TEST_SOURCE_QUICK,
+  TEST_SOURCE_ALL,
+} from "./pipeline_utils";
 import type { PipelineStep, Policy } from "./types";
+import * as complianceData from "../../data/compliancePrompts";
 
 const makeStep = (guardrail: string, overrides?: Partial<PipelineStep>): PipelineStep => ({
   guardrail,
@@ -183,3 +187,35 @@ describe("complianceMatchExpected", () => {
     expect(complianceMatchExpected("fail", "allow")).toBe(false);
   });
 });
+
+describe("getPromptsForTestSource", () => {
+  it("should return an empty array for quick_chat source", () => {
+    expect(getPromptsForTestSource(TEST_SOURCE_QUICK)).toEqual([]);
+  });
+
+  it("should return all compliance prompts for __all__ source", () => {
+    const mockPrompts = [
+      { id: "p1", framework: "fw", category: "c", categoryIcon: "", categoryDescription: "", prompt: "test", expectedResult: "pass" as const },
+    ];
+    vi.spyOn(complianceData, "getComplianceDatasetPrompts").mockReturnValue(mockPrompts);
+    const result = getPromptsForTestSource(TEST_SOURCE_ALL);
+    expect(result).toEqual(mockPrompts);
+    vi.restoreAllMocks();
+  });
+
+  it("should return prompts for a matching framework name", () => {
+    const prompt = { id: "p1", framework: "GDPR", category: "c", categoryIcon: "", categoryDescription: "", prompt: "test", expectedResult: "fail" as const };
+    vi.spyOn(complianceData, "getFrameworks").mockReturnValue([
+      { name: "GDPR", icon: "", description: "", categories: [{ name: "c", icon: "", description: "", prompts: [prompt] }] },
+    ]);
+    const result = getPromptsForTestSource("GDPR");
+    expect(result).toEqual([prompt]);
+    vi.restoreAllMocks();
+  });
+
+  it("should return an empty array for an unrecognized source", () => {
+    vi.spyOn(complianceData, "getFrameworks").mockReturnValue([]);
+    expect(getPromptsForTestSource("nonexistent")).toEqual([]);
+    vi.restoreAllMocks();
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_icons.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_utils.tsx
similarity index 99%
rename from ui/litellm-dashboard/src/components/policies/pipeline_icons.tsx
rename to ui/litellm-dashboard/src/components/policies/pipeline_utils.tsx
index b2de8d78b52..44c59579bbc 100644
--- a/ui/litellm-dashboard/src/components/policies/pipeline_icons.tsx
+++ b/ui/litellm-dashboard/src/components/policies/pipeline_utils.tsx
@@ -156,6 +156,8 @@ export const PassIcon: React.FC = () => (
 export const FailIcon: React.FC = () => (
   <svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="#f87171" strokeWidth="2.5" strokeLinecap="round" strokeLinejoin="round" style={{ flexShrink: 0 }}>
     <circle cx="12" cy="12" r="10" />
+    <path d="M15 9l-6 6" />
+    <path d="M9 9l6 6" />
   </svg>
 );