mirror of
https://github.com/tiennm99/litellm.git
synced 2026-06-18 00:48:01 +00:00
Ishaan Jaffer
289 lines
9.1 KiB
Python
289 lines
9.1 KiB
Python
import os
|
|
import sys
|
|
from datetime import datetime
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import polars as pl
|
|
import pytest
|
|
|
|
sys.path.insert(0, os.path.abspath("../../../.."))
|
|
|
|
from litellm.integrations.cloudzero.transform import CBFTransformer
|
|
from litellm.types.integrations.cloudzero import CBFRecord
|
|
|
|
|
|
class TestCBFTransformer:
|
|
"""Test suite for CBFTransformer class."""
|
|
|
|
def test_init(self):
|
|
"""Test CBFTransformer initialization."""
|
|
transformer = CBFTransformer()
|
|
assert hasattr(transformer, "czrn_generator")
|
|
assert transformer.czrn_generator is not None
|
|
|
|
def test_transform_empty_dataframe(self):
|
|
"""Test transform method with empty DataFrame."""
|
|
transformer = CBFTransformer()
|
|
empty_df = pl.DataFrame()
|
|
|
|
result = transformer.transform(empty_df)
|
|
|
|
assert result.is_empty()
|
|
assert isinstance(result, pl.DataFrame)
|
|
|
|
def test_transform_with_zero_successful_requests(self):
|
|
"""Test transform method filters out records with zero successful_requests."""
|
|
transformer = CBFTransformer()
|
|
data = pl.DataFrame(
|
|
{
|
|
"date": ["2025-01-19"],
|
|
"successful_requests": [0],
|
|
"spend": [10.0],
|
|
"entity_id": ["test_entity"],
|
|
"model": ["gpt-4"],
|
|
}
|
|
)
|
|
|
|
result = transformer.transform(data)
|
|
|
|
assert result.is_empty()
|
|
|
|
def test_transform_with_valid_data(self):
|
|
"""Test transform method with valid data."""
|
|
transformer = CBFTransformer()
|
|
with patch.object(transformer, "_create_cbf_record") as mock_create:
|
|
mock_create.return_value = CBFRecord({"test": "data"})
|
|
|
|
data = pl.DataFrame(
|
|
{
|
|
"date": ["2025-01-19"],
|
|
"successful_requests": [5],
|
|
"spend": [10.0],
|
|
"entity_id": ["test_entity"],
|
|
"model": ["gpt-4"],
|
|
}
|
|
)
|
|
|
|
result = transformer.transform(data)
|
|
|
|
assert len(result) == 1
|
|
mock_create.assert_called_once()
|
|
|
|
def test_transform_handles_czrn_generation_failures(self):
|
|
"""Test transform method handles CZRN generation failures gracefully."""
|
|
transformer = CBFTransformer()
|
|
with patch.object(transformer, "_create_cbf_record") as mock_create:
|
|
mock_create.side_effect = Exception("CZRN generation failed")
|
|
|
|
data = pl.DataFrame(
|
|
{
|
|
"date": ["2025-01-19"],
|
|
"successful_requests": [5],
|
|
"spend": [10.0],
|
|
"entity_id": ["test_entity"],
|
|
"model": ["gpt-4"],
|
|
}
|
|
)
|
|
|
|
result = transformer.transform(data)
|
|
|
|
assert result.is_empty()
|
|
|
|
def test_create_cbf_record(self):
|
|
"""Test _create_cbf_record method with valid row data."""
|
|
transformer = CBFTransformer()
|
|
with (
|
|
patch.object(
|
|
transformer.czrn_generator, "create_from_litellm_data"
|
|
) as mock_czrn,
|
|
patch.object(
|
|
transformer.czrn_generator, "extract_components"
|
|
) as mock_extract,
|
|
):
|
|
|
|
mock_czrn.return_value = "test-czrn"
|
|
mock_extract.return_value = (
|
|
"service",
|
|
"provider",
|
|
"region",
|
|
"account",
|
|
"resource",
|
|
"local_id",
|
|
)
|
|
|
|
row = {
|
|
"date": "2025-01-19",
|
|
"spend": 10.5,
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 50,
|
|
"entity_id": "test_entity",
|
|
"model": "gpt-4",
|
|
"entity_type": "user",
|
|
"model_group": "openai",
|
|
"custom_llm_provider": "openai",
|
|
"api_key": "sk-test123",
|
|
"api_requests": 5,
|
|
"successful_requests": 5,
|
|
"failed_requests": 0,
|
|
}
|
|
|
|
result = transformer._create_cbf_record(row)
|
|
|
|
assert isinstance(result, CBFRecord)
|
|
assert result["cost/cost"] == 10.5
|
|
assert result["usage/amount"] == 150 # 100 + 50
|
|
assert result["usage/units"] == "tokens"
|
|
assert result["resource/id"] == "test-czrn"
|
|
|
|
def test_create_cbf_record_adds_user_email_tag(self):
|
|
"""Test that user_email field is emitted as a resource tag when present."""
|
|
transformer = CBFTransformer()
|
|
with (
|
|
patch.object(
|
|
transformer.czrn_generator, "create_from_litellm_data"
|
|
) as mock_czrn,
|
|
patch.object(
|
|
transformer.czrn_generator, "extract_components"
|
|
) as mock_extract,
|
|
):
|
|
|
|
mock_czrn.return_value = "test-czrn"
|
|
mock_extract.return_value = (
|
|
"service",
|
|
"provider",
|
|
"region",
|
|
"account",
|
|
"resource",
|
|
"local_id",
|
|
)
|
|
|
|
row = {
|
|
"date": "2025-01-19",
|
|
"spend": 1.0,
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 5,
|
|
"model": "gpt-4",
|
|
"api_key": "sk-useremail",
|
|
"team_id": "team-123",
|
|
"team_alias": "Dev Team",
|
|
"user_email": "user@example.com",
|
|
}
|
|
|
|
result = transformer._create_cbf_record(row)
|
|
|
|
assert result["resource/tag:user_email"] == "user@example.com"
|
|
|
|
def test_create_cbf_record_omits_empty_user_email(self):
|
|
"""Test that empty user_email values are not added as resource tags."""
|
|
transformer = CBFTransformer()
|
|
with (
|
|
patch.object(
|
|
transformer.czrn_generator, "create_from_litellm_data"
|
|
) as mock_czrn,
|
|
patch.object(
|
|
transformer.czrn_generator, "extract_components"
|
|
) as mock_extract,
|
|
):
|
|
|
|
mock_czrn.return_value = "test-czrn"
|
|
mock_extract.return_value = (
|
|
"service",
|
|
"provider",
|
|
"region",
|
|
"account",
|
|
"resource",
|
|
"local_id",
|
|
)
|
|
|
|
row = {
|
|
"date": "2025-01-19",
|
|
"spend": 1.0,
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 5,
|
|
"model": "gpt-4",
|
|
"api_key": "sk-useremail",
|
|
"team_id": "team-123",
|
|
"team_alias": "Dev Team",
|
|
"user_email": None,
|
|
}
|
|
|
|
result = transformer._create_cbf_record(row)
|
|
|
|
assert "resource/tag:user_email" not in result
|
|
|
|
def test_create_cbf_record_minimal_data(self):
|
|
"""Test _create_cbf_record method with minimal row data."""
|
|
transformer = CBFTransformer()
|
|
with (
|
|
patch.object(
|
|
transformer.czrn_generator, "create_from_litellm_data"
|
|
) as mock_czrn,
|
|
patch.object(
|
|
transformer.czrn_generator, "extract_components"
|
|
) as mock_extract,
|
|
):
|
|
|
|
mock_czrn.return_value = "test-czrn"
|
|
mock_extract.return_value = (
|
|
"service",
|
|
"provider",
|
|
"region",
|
|
"account",
|
|
"resource",
|
|
"local_id",
|
|
)
|
|
|
|
row = {"date": "2025-01-19", "spend": 0.0}
|
|
|
|
result = transformer._create_cbf_record(row)
|
|
|
|
assert isinstance(result, CBFRecord)
|
|
assert result["cost/cost"] == 0.0
|
|
assert result["usage/amount"] == 0 # no tokens
|
|
assert result["usage/units"] == "tokens"
|
|
|
|
def test_parse_date_with_valid_string(self):
|
|
"""Test _parse_date method with valid date string."""
|
|
transformer = CBFTransformer()
|
|
|
|
result = transformer._parse_date("2025-01-19")
|
|
|
|
assert isinstance(result, datetime)
|
|
assert result.year == 2025
|
|
assert result.month == 1
|
|
assert result.day == 19
|
|
|
|
def test_parse_date_with_datetime_object(self):
|
|
"""Test _parse_date method with datetime object."""
|
|
transformer = CBFTransformer()
|
|
dt = datetime(2025, 1, 19)
|
|
|
|
result = transformer._parse_date(dt)
|
|
|
|
assert result == dt
|
|
|
|
def test_parse_date_with_none(self):
|
|
"""Test _parse_date method with None."""
|
|
transformer = CBFTransformer()
|
|
|
|
result = transformer._parse_date(None)
|
|
|
|
assert result is None
|
|
|
|
def test_parse_date_with_invalid_string(self):
|
|
"""Test _parse_date method with invalid date string."""
|
|
transformer = CBFTransformer()
|
|
|
|
result = transformer._parse_date("invalid-date")
|
|
|
|
assert result is None
|
|
|
|
def test_parse_date_with_iso_format(self):
|
|
"""Test _parse_date method with ISO format string."""
|
|
transformer = CBFTransformer()
|
|
|
|
result = transformer._parse_date("2025-01-19T10:30:00Z")
|
|
|
|
assert isinstance(result, datetime)
|
|
assert result.year == 2025
|