diff --git a/litellm/proxy/spend_tracking/cloudzero_endpoints.py b/litellm/proxy/spend_tracking/cloudzero_endpoints.py index 7aecf7a645..2cf4ce8f16 100644 --- a/litellm/proxy/spend_tracking/cloudzero_endpoints.py +++ b/litellm/proxy/spend_tracking/cloudzero_endpoints.py @@ -296,7 +296,7 @@ def is_cloudzero_setup_in_config() -> bool: async def is_cloudzero_setup() -> bool: """ - Check if CloudZero is setup in either config.yaml/env vars OR database + Check if CloudZero is setup in either config.yaml/env vars OR database. CloudZero is considered setup if: - CloudZero is configured in config.yaml callbacks, OR @@ -500,3 +500,70 @@ async def cloudzero_export( status_code=500, detail={"error": f"Failed to perform CloudZero export: {str(e)}"}, ) + + +@router.delete( + "/cloudzero/delete", + tags=["CloudZero"], + dependencies=[Depends(user_api_key_auth)], + response_model=CloudZeroInitResponse, +) +async def delete_cloudzero_settings( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Delete CloudZero settings from the database. + + This endpoint removes the CloudZero configuration (API key, connection ID, timezone) + from the proxy database. Only the CloudZero settings entry will be deleted; + other configuration values in the database will remain unchanged. + + Only admin users can delete CloudZero settings. + """ + # Validation + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={"error": CommonProxyErrors.not_allowed_access.value}, + ) + + try: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + # Check if CloudZero settings exist + cloudzero_config = await prisma_client.db.litellm_config.find_first( + where={"param_name": "cloudzero_settings"} + ) + + if cloudzero_config is None: + raise HTTPException( + status_code=404, + detail={"error": "CloudZero settings not found"}, + ) + + # Delete only the CloudZero settings entry + # This uses a specific where clause to target only the cloudzero_settings row + await prisma_client.db.litellm_config.delete( + where={"param_name": "cloudzero_settings"} + ) + + verbose_proxy_logger.info("CloudZero settings deleted successfully") + + return CloudZeroInitResponse( + message="CloudZero settings deleted successfully", status="success" + ) + + except HTTPException as e: + raise e + except Exception as e: + verbose_proxy_logger.error(f"Error deleting CloudZero settings: {str(e)}") + raise HTTPException( + status_code=500, + detail={"error": f"Failed to delete CloudZero settings: {str(e)}"}, + ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_cloudzero_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_cloudzero_endpoints.py new file mode 100644 index 0000000000..8ff5774bf5 --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_cloudzero_endpoints.py @@ -0,0 +1,79 @@ +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../../..") +) + +import litellm.proxy.proxy_server as ps +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.proxy_server import app + + +@pytest.fixture +def client(): + return TestClient(app) + + +@pytest.mark.asyncio +async def test_delete_cloudzero_settings_success(client, monkeypatch): + mock_config = MagicMock() + mock_config.param_name = "cloudzero_settings" + mock_config.param_value = {"api_key": "encrypted_key", "connection_id": "conn_123", "timezone": "UTC"} + + mock_litellm_config = MagicMock() + mock_litellm_config.find_first = AsyncMock(return_value=mock_config) + mock_litellm_config.delete = AsyncMock(return_value=mock_config) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_config = mock_litellm_config + + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + try: + response = client.delete("/cloudzero/delete") + assert response.status_code == 200 + data = response.json() + assert data["message"] == "CloudZero settings deleted successfully" + assert data["status"] == "success" + mock_litellm_config.find_first.assert_awaited_once() + mock_litellm_config.delete.assert_awaited_once() + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_delete_cloudzero_settings_not_found(client, monkeypatch): + mock_litellm_config = MagicMock() + mock_litellm_config.find_first = AsyncMock(return_value=None) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_config = mock_litellm_config + + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + try: + response = client.delete("/cloudzero/delete") + assert response.status_code == 404 + data = response.json() + assert "error" in data["detail"] + assert "CloudZero settings not found" in data["detail"]["error"] + mock_litellm_config.find_first.assert_awaited_once() + mock_litellm_config.delete.assert_not_called() + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) +