mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-10 21:04:46 +00:00
Ishaan Jaffer
141 lines
5.1 KiB
Python
141 lines
5.1 KiB
Python
"""
|
|
Test the CloudZero dry run endpoint functionality
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import polars as pl
|
|
import pytest
|
|
|
|
sys.path.insert(0, os.path.abspath("../../../.."))
|
|
|
|
from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
|
|
|
|
|
|
class TestCloudZeroDryRunEndpoint:
|
|
"""Test suite for CloudZero dry run endpoint functionality."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dry_run_export_usage_data_returns_data(self):
|
|
"""
|
|
Test that dry_run_export_usage_data returns expected data structure
|
|
instead of just logging to console.
|
|
"""
|
|
logger = CloudZeroLogger()
|
|
|
|
# Mock database data
|
|
mock_usage_data = pl.DataFrame(
|
|
{
|
|
"date": ["2025-01-19", "2025-01-20"],
|
|
"model": ["gpt-4", "gpt-3.5-turbo"],
|
|
"custom_llm_provider": ["openai", "openai"],
|
|
"team_id": ["team1", "team2"],
|
|
"team_alias": ["Team One", "Team Two"],
|
|
"api_key_alias": ["key1", "key2"],
|
|
"user_email": ["one@example.com", None],
|
|
"prompt_tokens": [100, 200],
|
|
"completion_tokens": [50, 100],
|
|
"spend": [0.01, 0.02],
|
|
"successful_requests": [1, 2],
|
|
}
|
|
)
|
|
|
|
# Mock CBF transformed data
|
|
mock_cbf_data = pl.DataFrame(
|
|
{
|
|
"time/usage_start": ["2025-01-19T00:00:00Z", "2025-01-20T00:00:00Z"],
|
|
"cost/cost": [0.01, 0.02],
|
|
"usage/amount": [150, 300],
|
|
"resource/service": ["openai", "openai"],
|
|
"resource/account": ["litellm", "litellm"],
|
|
"resource/region": ["us-east-1", "us-east-1"],
|
|
"resource/id": ["gpt-4", "gpt-3.5-turbo"],
|
|
"entity_type": ["user", "user"],
|
|
"entity_id": ["team1", "team2"],
|
|
"resource/tag:team_id": ["team1", "team2"],
|
|
"resource/tag:team_alias": ["Team One", "Team Two"],
|
|
"resource/tag:api_key_alias": ["key1", "key2"],
|
|
"resource/tag:user_email": ["one@example.com", "N/A"],
|
|
}
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.integrations.cloudzero.database.LiteLLMDatabase"
|
|
) as mock_db_class,
|
|
patch(
|
|
"litellm.integrations.cloudzero.transform.CBFTransformer"
|
|
) as mock_transformer_class,
|
|
):
|
|
|
|
# Setup mocks
|
|
mock_db = AsyncMock()
|
|
mock_db.get_usage_data.return_value = mock_usage_data
|
|
mock_db_class.return_value = mock_db
|
|
|
|
mock_transformer = MagicMock()
|
|
mock_transformer.transform.return_value = mock_cbf_data
|
|
mock_transformer_class.return_value = mock_transformer
|
|
|
|
# Call the method
|
|
result = await logger.dry_run_export_usage_data(limit=1000)
|
|
|
|
# Verify the result structure
|
|
assert isinstance(result, dict)
|
|
assert "usage_data" in result
|
|
assert "cbf_data" in result
|
|
assert "summary" in result
|
|
|
|
# Verify usage_data
|
|
assert isinstance(result["usage_data"], list)
|
|
assert len(result["usage_data"]) == 2
|
|
assert result["usage_data"][0]["model"] == "gpt-4"
|
|
assert result["usage_data"][1]["model"] == "gpt-3.5-turbo"
|
|
|
|
# Verify cbf_data
|
|
assert isinstance(result["cbf_data"], list)
|
|
assert len(result["cbf_data"]) == 2
|
|
assert result["cbf_data"][0]["cost/cost"] == 0.01
|
|
assert result["cbf_data"][1]["cost/cost"] == 0.02
|
|
assert result["cbf_data"][0]["resource/tag:user_email"] == "one@example.com"
|
|
|
|
# Verify summary
|
|
summary = result["summary"]
|
|
assert summary["total_records"] == 2
|
|
assert summary["total_cost"] == 0.03
|
|
assert summary["total_tokens"] == 450 # 150 + 300
|
|
assert summary["unique_accounts"] == 1
|
|
assert summary["unique_services"] == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dry_run_export_usage_data_empty_data(self):
|
|
"""
|
|
Test that dry_run_export_usage_data handles empty data gracefully.
|
|
"""
|
|
logger = CloudZeroLogger()
|
|
|
|
# Mock empty database data
|
|
mock_empty_data = pl.DataFrame()
|
|
|
|
with patch(
|
|
"litellm.integrations.cloudzero.database.LiteLLMDatabase"
|
|
) as mock_db_class:
|
|
|
|
# Setup mocks
|
|
mock_db = AsyncMock()
|
|
mock_db.get_usage_data.return_value = mock_empty_data
|
|
mock_db_class.return_value = mock_db
|
|
|
|
# Call the method
|
|
result = await logger.dry_run_export_usage_data(limit=1000)
|
|
|
|
# Verify the result structure for empty data
|
|
assert isinstance(result, dict)
|
|
assert result["usage_data"] == []
|
|
assert result["cbf_data"] == []
|
|
assert result["summary"]["total_records"] == 0
|
|
assert result["summary"]["total_cost"] == 0
|
|
assert result["summary"]["total_tokens"] == 0
|