From 1c611c4d6b98e9c18365b640c1cbbea6fd7650ee Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 23 Mar 2026 22:14:09 -0700 Subject: [PATCH] Add try/finally to compliance-dataset path and test coverage - Wrap compliance-dataset loop in try/finally so isRunning resets consistently with the quick-chat branch - Add 4 tests covering the compliance-dataset execution path: matched/total count, mismatch indicator, per-prompt errors, and call count verification Co-Authored-By: Claude Opus 4.6 (1M context) --- .../policies/pipeline_test_drawer.test.tsx | 168 +++++++++++++++++- .../policies/pipeline_test_drawer.tsx | 37 ++-- 2 files changed, 186 insertions(+), 19 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 f2a4cca0389..be5a09b4387 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 @@ -1,14 +1,68 @@ -import { screen, waitFor } from "@testing-library/react"; +import { screen, waitFor, fireEvent, act } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "../../../tests/test-utils"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { PipelineTestPanel } from "./pipeline_test_drawer"; import type { GuardrailPipeline, PipelineTestResult } from "./types"; +import type { CompliancePrompt } from "../../data/compliancePrompts"; import * as networking from "../networking"; +const mockPassPrompt: CompliancePrompt = { + id: "cp-1", + framework: "TestFW", + category: "cat", + categoryIcon: "", + categoryDescription: "", + prompt: "Is this safe?", + expectedResult: "pass", +}; + +const mockFailPrompt: CompliancePrompt = { + id: "cp-2", + framework: "TestFW", + category: "cat", + categoryIcon: "", + categoryDescription: "", + prompt: "How to hack a server?", + expectedResult: "fail", +}; + vi.mock("../networking"); vi.mock("../../data/compliancePrompts", () => ({ - getFrameworks: () => [], + getFrameworks: () => [ + { + name: "TestFW", + icon: "", + description: "", + categories: [ + { + name: "cat", + icon: "", + description: "", + prompts: [ + { + id: "cp-1", + framework: "TestFW", + category: "cat", + categoryIcon: "", + categoryDescription: "", + prompt: "Is this safe?", + expectedResult: "pass", + }, + { + id: "cp-2", + framework: "TestFW", + category: "cat", + categoryIcon: "", + categoryDescription: "", + prompt: "How to hack a server?", + expectedResult: "fail", + }, + ], + }, + ], + }, + ], getComplianceDatasetPrompts: () => [], })); @@ -61,6 +115,16 @@ const defaultProps = { onClose: vi.fn(), }; +async function selectAntOption(user: ReturnType, title: string) { + const selector = document.querySelector(".ant-select-selector") as HTMLElement; + 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; + await user.click(option); +} + describe("PipelineTestPanel", () => { beforeEach(() => { vi.clearAllMocks(); @@ -195,3 +259,103 @@ describe("PipelineTestPanel", () => { }); }); }); + +describe("PipelineTestPanel compliance-dataset mode", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should show matched/total count after running a compliance dataset", async () => { + const allowResult: PipelineTestResult = { + terminal_action: "allow", + step_results: [], + modified_data: null, + error_message: null, + modify_response_message: null, + }; + const blockResult: PipelineTestResult = { + terminal_action: "block", + step_results: [], + modified_data: null, + error_message: null, + modify_response_message: null, + }; + vi.mocked(networking.testPipelineCall) + .mockResolvedValueOnce(allowResult) + .mockResolvedValueOnce(blockResult); + + const user = userEvent.setup(); + renderWithProviders(); + await selectAntOption(user, "TestFW"); + await user.click(screen.getByRole("button", { name: /run test/i })); + + await waitFor(() => { + expect(screen.getByText(/2 \/ 2 matched/)).toBeInTheDocument(); + }); + }); + + it("should show a mismatch indicator when actual result differs from expected", async () => { + const blockResult: PipelineTestResult = { + terminal_action: "block", + step_results: [], + modified_data: null, + error_message: null, + modify_response_message: null, + }; + // Both prompts return "block": pass-expected prompt mismatches, fail-expected matches + vi.mocked(networking.testPipelineCall) + .mockResolvedValueOnce(blockResult) + .mockResolvedValueOnce(blockResult); + + const user = userEvent.setup(); + renderWithProviders(); + await selectAntOption(user, "TestFW"); + await user.click(screen.getByRole("button", { name: /run test/i })); + + await waitFor(() => { + expect(screen.getByText(/1 \/ 2 matched/)).toBeInTheDocument(); + }); + }); + + it("should display per-prompt error when an API call fails in dataset mode", async () => { + vi.mocked(networking.testPipelineCall) + .mockRejectedValueOnce(new Error("Timeout on prompt 1")) + .mockResolvedValueOnce({ + terminal_action: "block", + step_results: [], + modified_data: null, + error_message: null, + modify_response_message: null, + }); + + const user = userEvent.setup(); + renderWithProviders(); + await selectAntOption(user, "TestFW"); + await user.click(screen.getByRole("button", { name: /run test/i })); + + await waitFor(() => { + expect(screen.getByText("Timeout on prompt 1")).toBeInTheDocument(); + }); + expect(screen.getByText(/1 \/ 2 matched/)).toBeInTheDocument(); + }); + + it("should call testPipelineCall once per prompt in the dataset", async () => { + const allowResult: PipelineTestResult = { + terminal_action: "allow", + step_results: [], + modified_data: null, + error_message: null, + modify_response_message: null, + }; + vi.mocked(networking.testPipelineCall).mockResolvedValue(allowResult); + + const user = userEvent.setup(); + renderWithProviders(); + await selectAntOption(user, "TestFW"); + await user.click(screen.getByRole("button", { name: /run test/i })); + + await waitFor(() => { + expect(networking.testPipelineCall).toHaveBeenCalledTimes(2); + }); + }); +}); 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 670ff55d480..2cb425f8552 100644 --- a/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.tsx +++ b/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.tsx @@ -107,25 +107,28 @@ export const PipelineTestPanel: React.FC = ({ } const entries: ComplianceRunEntry[] = []; - for (const prompt of promptsForSource) { - try { - const data = await testPipelineCall(accessToken, pipeline, [ - { role: "user", content: prompt.prompt }, - ]); - const matched = complianceMatchExpected(prompt.expectedResult, data.terminal_action); - entries.push({ prompt, result: data, matched }); - } catch (e) { - const errMsg = e instanceof Error ? e.message : String(e); - entries.push({ - prompt, - result: null, - error: errMsg, - matched: false, - }); + try { + for (const prompt of promptsForSource) { + try { + const data = await testPipelineCall(accessToken, pipeline, [ + { role: "user", content: prompt.prompt }, + ]); + const matched = complianceMatchExpected(prompt.expectedResult, data.terminal_action); + entries.push({ prompt, result: data, matched }); + } catch (e) { + const errMsg = e instanceof Error ? e.message : String(e); + entries.push({ + prompt, + result: null, + error: errMsg, + matched: false, + }); + } } + setComplianceResults(entries); + } finally { + setIsRunning(false); } - setComplianceResults(entries); - setIsRunning(false); }; return (