mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
285 lines
9 KiB
Python
285 lines
9 KiB
Python
from datetime import datetime
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import polars as pl
|
|
import pytest
|
|
|
|
|
|
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
|