diff --git a/ui/litellm-dashboard/src/components/policies/index.tsx b/ui/litellm-dashboard/src/components/policies/index.tsx index 10a6848bd52..19d4042266f 100644 --- a/ui/litellm-dashboard/src/components/policies/index.tsx +++ b/ui/litellm-dashboard/src/components/policies/index.tsx @@ -16,11 +16,11 @@ import PolicyTemplates from "./policy_templates"; import GuardrailSelectionModal from "./guardrail_selection_modal"; import TemplateParameterModal from "./template_parameter_modal"; import AiSuggestionModal from "./ai_suggestion_modal"; +import { useDeletePolicyAttachment } from "@/hooks/policies/useDeletePolicyAttachment"; import { getPoliciesList, deletePolicyCall, getPolicyAttachmentsList, - deletePolicyAttachmentCall, getGuardrailsList, getPolicyInfo, createPolicyCall, @@ -169,22 +169,11 @@ const PoliciesPanel: React.FC = ({ setPolicyToDelete(null); }; - const deleteAttachmentMutation = useMutation({ - mutationFn: async (attachmentId: string) => { - if (!accessToken) { - throw new Error("Access token is required"); - } - return deletePolicyAttachmentCall(accessToken, attachmentId); - }, - onSuccess: async () => { - MessageManager.success("Attachment deleted successfully"); - await fetchAttachments(); - }, - onError: (error) => { - console.error("Error deleting attachment:", error); - MessageManager.error("Failed to delete attachment"); - }, + const deleteAttachmentMutation = useDeletePolicyAttachment({ + accessToken, + onSuccess: fetchAttachments, }); + const handleDeleteAttachmentClick = (attachmentId: string) => { const attachment = attachmentsList.find((a) => a.attachment_id === attachmentId) || null; setAttachmentToDelete(attachment); @@ -196,14 +185,14 @@ const PoliciesPanel: React.FC = ({ setAttachmentToDelete(null); }; - const handleAttachmentDeleteConfirm = async () => { + const handleAttachmentDeleteConfirm = () => { if (!attachmentToDelete) return; - try { - await deleteAttachmentMutation.mutateAsync(attachmentToDelete.attachment_id); - } finally { - setIsDeleteAttachmentModalOpen(false); - setAttachmentToDelete(null); - } + deleteAttachmentMutation.mutate(attachmentToDelete.attachment_id, { + onSettled: () => { + setIsDeleteAttachmentModalOpen(false); + setAttachmentToDelete(null); + }, + }); }; const handleAttachmentSuccess = () => { diff --git a/ui/litellm-dashboard/src/hooks/policies/useDeletePolicyAttachment.test.ts b/ui/litellm-dashboard/src/hooks/policies/useDeletePolicyAttachment.test.ts new file mode 100644 index 00000000000..fe63286ed01 --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/policies/useDeletePolicyAttachment.test.ts @@ -0,0 +1,80 @@ +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { useDeletePolicyAttachment } from "./useDeletePolicyAttachment"; +import { deletePolicyAttachmentCall } from "@/components/networking"; +import MessageManager from "@/components/molecules/message_manager"; +import { vi, describe, beforeEach, it, expect } from "vitest"; + +// Mock dependencies +vi.mock("@/components/networking", () => ({ + deletePolicyAttachmentCall: vi.fn(), +})); + +vi.mock("@/components/molecules/message_manager", () => ({ + default: { + success: vi.fn(), + error: vi.fn(), + }, +})); + +describe("useDeletePolicyAttachment", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient(); + vi.clearAllMocks(); + }); + + const wrapper = ({ children }: { children: React.ReactNode }) => ( + {children} + ); + + it("should successfully delete a policy attachment and call onSuccess", async () => { + const mockOnSuccess = vi.fn(); + (deletePolicyAttachmentCall as any).mockResolvedValue({}); + + const { result } = renderHook( + () => + useDeletePolicyAttachment({ + accessToken: "test-token", + onSuccess: mockOnSuccess, + }), + { wrapper } + ); + + result.current.mutate("attachment-1"); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(deletePolicyAttachmentCall).toHaveBeenCalledWith("test-token", "attachment-1"); + expect(MessageManager.success).toHaveBeenCalledWith("Attachment deleted successfully"); + expect(mockOnSuccess).toHaveBeenCalled(); + }); + + it("should handle error when deleting policy attachment", async () => { + const mockOnError = vi.fn(); + const error = new Error("Delete failed"); + (deletePolicyAttachmentCall as any).mockRejectedValue(error); + + const { result } = renderHook( + () => + useDeletePolicyAttachment({ + accessToken: "test-token", + onError: mockOnError, + }), + { wrapper } + ); + + result.current.mutate("attachment-1"); + + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(deletePolicyAttachmentCall).toHaveBeenCalledWith("test-token", "attachment-1"); + expect(MessageManager.error).toHaveBeenCalledWith("Failed to delete attachment"); + expect(mockOnError).toHaveBeenCalledWith(error); + }); +}); diff --git a/ui/litellm-dashboard/src/hooks/policies/useDeletePolicyAttachment.ts b/ui/litellm-dashboard/src/hooks/policies/useDeletePolicyAttachment.ts new file mode 100644 index 00000000000..d20d2279f6e --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/policies/useDeletePolicyAttachment.ts @@ -0,0 +1,37 @@ +import { useMutation } from "@tanstack/react-query"; +import { deletePolicyAttachmentCall } from "@/components/networking"; +import MessageManager from "@/components/molecules/message_manager"; + +interface UseDeletePolicyAttachmentProps { + accessToken: string | null; + onSuccess?: () => void; + onError?: (error: any) => void; +} + +export const useDeletePolicyAttachment = ({ + accessToken, + onSuccess, + onError, +}: UseDeletePolicyAttachmentProps) => { + return useMutation({ + mutationFn: async (attachmentId: string) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return deletePolicyAttachmentCall(accessToken, attachmentId); + }, + onSuccess: () => { + MessageManager.success("Attachment deleted successfully"); + if (onSuccess) { + onSuccess(); + } + }, + onError: (error) => { + console.error("Error deleting attachment:", error); + MessageManager.error("Failed to delete attachment"); + if (onError) { + onError(error); + } + }, + }); +};