refactor: migrate policy attachment deletion to useMutation hook with tests

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<PoliciesPanelProps> = ({
     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<PoliciesPanelProps> = ({
     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 = () => {

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 }) => (
+    <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+  );
+
+  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);
+  });
+});

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);
+      }
+    },
+  });
+};