test_dry_run_export_usage_data_cbf_transformation_failure

This commit is contained in:
Ishaan Jaff 2025-09-06 13:59:28 -07:00
parent 7dce6a81a8
commit fc7accad74

View file

@ -123,41 +123,3 @@ class TestCloudZeroDryRunEndpoint:
assert result['summary']['total_cost'] == 0
assert result['summary']['total_tokens'] == 0
@pytest.mark.asyncio
async def test_dry_run_export_usage_data_cbf_transformation_failure(self):
"""
Test that dry_run_export_usage_data handles CBF transformation failure gracefully.
"""
logger = CloudZeroLogger()
# Mock database data
mock_usage_data = pl.DataFrame({
'date': ['2025-01-19'],
'model': ['gpt-4'],
'spend': [0.01],
'successful_requests': [1]
})
# Mock empty CBF data (transformation failed)
mock_empty_cbf_data = pl.DataFrame()
with patch('litellm.integrations.cloudzero.cloudzero.LiteLLMDatabase') as mock_db_class, \
patch('litellm.integrations.cloudzero.cloudzero.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_empty_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 handles CBF transformation failure
assert isinstance(result, dict)
assert len(result['usage_data']) == 1 # Usage data should still be present
assert result['cbf_data'] == [] # CBF data should be empty
assert result['summary']['total_cost'] == 0.01 # Should calculate from usage data