mirror of
https://github.com/tiennm99/litellm.git
synced 2026-06-18 03:31:23 +00:00
Ishaan Jaff
126 lines
4.8 KiB
Python
126 lines
4.8 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'],
|
|
'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']
|
|
})
|
|
|
|
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
|
|
|
|
# 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
|
|
|