litellm/tests/test_litellm/proxy/spend_tracking/test_cloudzero_endpoints.py

import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch

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)


@pytest.mark.asyncio
async def test_get_cloudzero_settings_success(client, monkeypatch):
    """Test GET /cloudzero/settings returns settings when configured"""
    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_prisma = MagicMock()
    mock_prisma.db = MagicMock()
    mock_prisma.db.litellm_config = mock_litellm_config

    monkeypatch.setattr(ps, "prisma_client", mock_prisma)

    # Mock the decrypt function to return a decrypted key
    with patch(
        "litellm.proxy.spend_tracking.cloudzero_endpoints.decrypt_value_helper"
    ) as mock_decrypt:
        mock_decrypt.return_value = "decrypted_api_key"

        # Mock the masker
        with patch(
            "litellm.proxy.spend_tracking.cloudzero_endpoints._sensitive_masker"
        ) as mock_masker:
            mock_masker.mask_dict.return_value = {"api_key": "test****key"}

            app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
                user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
            )

            try:
                response = client.get("/cloudzero/settings")
                assert response.status_code == 200
                data = response.json()
                assert data["connection_id"] == "conn_123"
                assert data["timezone"] == "UTC"
                assert data["status"] == "configured"
                assert data["api_key_masked"] == "test****key"
                mock_litellm_config.find_first.assert_awaited_once()
            finally:
                app.dependency_overrides.pop(ps.user_api_key_auth, None)


@pytest.mark.asyncio
async def test_get_cloudzero_settings_not_configured(client, monkeypatch):
    """Test GET /cloudzero/settings returns 200 with null values when not configured (consistent with other endpoints)"""
    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.get("/cloudzero/settings")
        # Should return 200 with null values (not 404) - consistent with other settings endpoints
        assert response.status_code == 200
        data = response.json()
        assert data["api_key_masked"] is None
        assert data["connection_id"] is None
        assert data["timezone"] is None
        assert data["status"] is None
        mock_litellm_config.find_first.assert_awaited_once()
    finally:
        app.dependency_overrides.pop(ps.user_api_key_auth, None)


@pytest.mark.asyncio
async def test_get_cloudzero_settings_empty_param_value(client, monkeypatch):
    """Test GET /cloudzero/settings returns 200 with null values when param_value is None"""
    mock_config = MagicMock()
    mock_config.param_name = "cloudzero_settings"
    mock_config.param_value = None

    mock_litellm_config = MagicMock()
    mock_litellm_config.find_first = 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.get("/cloudzero/settings")
        # Should return 200 with null values (not 404) - consistent with other settings endpoints
        assert response.status_code == 200
        data = response.json()
        assert data["api_key_masked"] is None
        assert data["connection_id"] is None
        assert data["timezone"] is None
        assert data["status"] is None
        mock_litellm_config.find_first.assert_awaited_once()
    finally:
        app.dependency_overrides.pop(ps.user_api_key_auth, None)