diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.integration.test.tsx index 969844c58ed..eb8692357fe 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.integration.test.tsx @@ -1,7 +1,35 @@ +import { useState } from "react"; import { fireEvent, render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; import PiiConfiguration from "./pii_configuration"; +const StatefulPiiConfiguration = () => { + const [selectedEntities, setSelectedEntities] = useState([]); + const [selectedActions, setSelectedActions] = useState<{ [key: string]: string }>({}); + + return ( + { + setSelectedEntities((previous) => + previous.includes(entity) ? previous.filter((selected) => selected !== entity) : [...previous, entity], + ); + }} + onActionSelect={(entity, action) => { + setSelectedActions((previous) => ({ ...previous, [entity]: action })); + }} + entityCategories={[ + { category: "Personal", entities: ["PERSON"] }, + { category: "Contact", entities: ["EMAIL"] }, + ]} + /> + ); +}; + describe("PiiConfiguration", () => { it("should render", () => { render( @@ -56,6 +84,31 @@ describe("PiiConfiguration", () => { expect(onActionSelect).toHaveBeenCalledWith("SE_PERSONNUMMER", "MASK"); }); + it("shows a newly added custom entity while a category filter is active", async () => { + const user = userEvent.setup(); + render(); + + const categoryInput = screen.getByPlaceholderText("Select categories to filter by"); + await user.click(categoryInput); + if (categoryInput.getAttribute("aria-expanded") !== "true") { + categoryInput.focus(); + await user.keyboard("{Enter}"); + } + await user.click(await screen.findByRole("option", { name: "Personal" })); + await user.keyboard("{Escape}"); + + expect(screen.getByText("PERSON")).toBeVisible(); + expect(screen.queryByText("EMAIL")).not.toBeInTheDocument(); + + fireEvent.change(screen.getByRole("textbox", { name: "Custom entity name" }), { + target: { value: "se personnummer" }, + }); + fireEvent.click(screen.getByRole("button", { name: "Add entity" })); + + expect(await screen.findByText("SE PERSONNUMMER")).toBeVisible(); + expect(screen.getByText("PERSON")).toBeVisible(); + }); + it("rejects invalid custom entity names", () => { const onEntitySelect = vi.fn(); const onActionSelect = vi.fn(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx index e6efdf53670..609899aa95a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx @@ -77,6 +77,9 @@ const PiiConfiguration: React.FC = ({ return; } + if (selectedCategories.length > 0 && !selectedCategories.includes("Custom")) { + setSelectedCategories([...selectedCategories, "Custom"]); + } onEntitySelect(normalizedEntityName); onActionSelect(normalizedEntityName, "MASK"); };