mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-08 08:08:37 +00:00
243 lines
9.8 KiB
Python
243 lines
9.8 KiB
Python
import json
|
|
import zoneinfo
|
|
from datetime import datetime, timezone
|
|
from typing import Final
|
|
from unittest.mock import MagicMock, Mock, patch
|
|
|
|
import httpx
|
|
import polars as pl
|
|
import pytest
|
|
import respx
|
|
|
|
import litellm
|
|
from litellm.integrations.cloudzero.cz_stream_api import CloudZeroStreamer
|
|
|
|
|
|
class TestCloudZeroStreamer:
|
|
"""Test suite for CloudZeroStreamer class."""
|
|
|
|
def test_init_with_defaults(self):
|
|
"""Test CloudZeroStreamer initialization with default parameters."""
|
|
streamer = CloudZeroStreamer(
|
|
api_key="test-key", connection_id="test-connection"
|
|
)
|
|
|
|
assert streamer.api_key == "test-key"
|
|
assert streamer.connection_id == "test-connection"
|
|
assert streamer.base_url == "https://api.cloudzero.com"
|
|
assert streamer.user_timezone == timezone.utc
|
|
|
|
def test_init_with_valid_timezone(self):
|
|
"""Test CloudZeroStreamer initialization with valid timezone."""
|
|
streamer = CloudZeroStreamer(
|
|
api_key="test-key",
|
|
connection_id="test-connection",
|
|
user_timezone="America/New_York",
|
|
)
|
|
|
|
assert streamer.user_timezone == zoneinfo.ZoneInfo("America/New_York")
|
|
|
|
def test_send_batched_with_valid_data(self):
|
|
"""Test send_batched method with valid data."""
|
|
streamer = CloudZeroStreamer("test-key", "test-connection")
|
|
with (
|
|
patch.object(streamer, "_group_by_date") as mock_group,
|
|
patch.object(streamer, "_send_daily_batch") as mock_send,
|
|
):
|
|
|
|
mock_group.return_value = {
|
|
"2025-01-19": pl.DataFrame({"test": ["data1"]}),
|
|
"2025-01-20": pl.DataFrame({"test": ["data2"]}),
|
|
}
|
|
|
|
data = pl.DataFrame({"test": ["data"]})
|
|
streamer.send_batched(data, "replace_hourly")
|
|
|
|
assert mock_send.call_count == 2
|
|
|
|
def test_group_by_date_valid_data(self):
|
|
"""Test _group_by_date method with valid data."""
|
|
streamer = CloudZeroStreamer("test-key", "test-connection")
|
|
with patch.object(streamer, "_parse_and_convert_timestamp") as mock_parse:
|
|
mock_parse.return_value = datetime(
|
|
2025, 1, 19, 10, 30, 0, tzinfo=timezone.utc
|
|
)
|
|
|
|
data = pl.DataFrame(
|
|
{"time/usage_start": ["2025-01-19T10:30:00Z"], "cost": [10.0]}
|
|
)
|
|
|
|
result = streamer._group_by_date(data)
|
|
|
|
assert "2025-01-19" in result
|
|
assert len(result["2025-01-19"]) == 1
|
|
|
|
def test_group_by_date_infers_schema_from_every_row(self):
|
|
"""Test daily batches retain optional string columns that are null for thousands of leading rows."""
|
|
streamer = CloudZeroStreamer("test-key", "test-connection")
|
|
leading_nulls = 10_000
|
|
rows = [
|
|
{"time/usage_start": "2025-01-19T10:30:00Z", "resource/tag:team_alias": None}
|
|
for _ in range(leading_nulls)
|
|
]
|
|
rows.append(
|
|
{"time/usage_start": "2025-01-19T10:30:00Z", "resource/tag:team_alias": "team-alias"}
|
|
)
|
|
data = pl.DataFrame(
|
|
rows,
|
|
schema={"time/usage_start": pl.String, "resource/tag:team_alias": pl.String},
|
|
)
|
|
|
|
result = streamer._group_by_date(data)
|
|
|
|
batch = result["2025-01-19"]
|
|
assert len(batch) == leading_nulls + 1
|
|
assert batch.schema["resource/tag:team_alias"] == pl.String
|
|
assert batch["resource/tag:team_alias"].null_count() == leading_nulls
|
|
assert batch.tail(1).item(0, "resource/tag:team_alias") == "team-alias"
|
|
|
|
def test_parse_and_convert_timestamp_utc(self):
|
|
"""Test _parse_and_convert_timestamp method with UTC timestamp."""
|
|
streamer = CloudZeroStreamer("test-key", "test-connection")
|
|
|
|
result = streamer._parse_and_convert_timestamp("2025-01-19T10:30:00Z")
|
|
|
|
assert result.year == 2025
|
|
assert result.month == 1
|
|
assert result.day == 19
|
|
assert result.hour == 10
|
|
assert result.minute == 30
|
|
assert result.tzinfo == timezone.utc
|
|
|
|
def test_parse_and_convert_timestamp_with_offset(self):
|
|
"""Test _parse_and_convert_timestamp method with timezone offset."""
|
|
streamer = CloudZeroStreamer("test-key", "test-connection")
|
|
|
|
result = streamer._parse_and_convert_timestamp("2025-01-19T10:30:00+05:00")
|
|
|
|
assert result.tzinfo == timezone.utc
|
|
assert result.hour == 5 # Converted to UTC
|
|
|
|
def test_parse_and_convert_timestamp_no_timezone(self):
|
|
"""Test _parse_and_convert_timestamp method without timezone info."""
|
|
streamer = CloudZeroStreamer(
|
|
"test-key", "test-connection", user_timezone="America/New_York"
|
|
)
|
|
|
|
result = streamer._parse_and_convert_timestamp("2025-01-19T10:30:00")
|
|
|
|
assert result.tzinfo == timezone.utc
|
|
|
|
def test_parse_and_convert_timestamp_invalid(self):
|
|
"""Test _parse_and_convert_timestamp method with invalid timestamp."""
|
|
streamer = CloudZeroStreamer("test-key", "test-connection")
|
|
|
|
with pytest.raises(ValueError, match="Could not parse timestamp 'invalid-timestamp': Invalid"):
|
|
streamer._parse_and_convert_timestamp("invalid-timestamp")
|
|
|
|
def test_prepare_batch_payload(self):
|
|
"""Test _prepare_batch_payload method."""
|
|
streamer = CloudZeroStreamer("test-key", "test-connection")
|
|
with patch.object(streamer, "_convert_cbf_to_api_format") as mock_convert:
|
|
mock_convert.return_value = {"test": "record"}
|
|
|
|
batch_data = pl.DataFrame({"cost": [10.0]})
|
|
result = streamer._prepare_batch_payload(
|
|
"2025-01-19", batch_data, "replace_hourly"
|
|
)
|
|
|
|
assert result["month"] == "2025-01"
|
|
assert result["operation"] == "replace_hourly"
|
|
assert len(result["data"]) == 1
|
|
|
|
def test_convert_cbf_to_api_format_valid_data(self):
|
|
"""Test _convert_cbf_to_api_format method with valid data."""
|
|
streamer = CloudZeroStreamer("test-key", "test-connection")
|
|
with patch.object(streamer, "_ensure_utc_timestamp") as mock_ensure:
|
|
mock_ensure.return_value = "2025-01-19T10:30:00Z"
|
|
|
|
row = {
|
|
"time/usage_start": "2025-01-19T10:30:00Z",
|
|
"cost/cost": 10.5,
|
|
"tokens": 100,
|
|
"text_field": "test",
|
|
}
|
|
|
|
result = streamer._convert_cbf_to_api_format(row)
|
|
|
|
assert result["cost/cost"] == "10.5"
|
|
assert result["tokens"] == "100"
|
|
assert result["text_field"] == "test"
|
|
|
|
def test_convert_cbf_to_api_format_float_precision(self):
|
|
"""Test _convert_cbf_to_api_format method handles float precision correctly."""
|
|
streamer = CloudZeroStreamer("test-key", "test-connection")
|
|
|
|
row = {"cost": 10.123456789012345, "large_float": 1234567890.0}
|
|
|
|
result = streamer._convert_cbf_to_api_format(row)
|
|
|
|
# Should avoid scientific notation
|
|
assert "e" not in result["cost"].lower()
|
|
assert "e" not in result["large_float"].lower()
|
|
|
|
|
|
@pytest.mark.parametrize("redact", (False, True))
|
|
def test_cloudzero_http_error_diagnostics_respect_privacy_without_changing_exception(
|
|
redact: bool, capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", redact)
|
|
marker: Final = "synthetic-private-upstream-content"
|
|
streamer: Final = CloudZeroStreamer("test-key", "test-connection")
|
|
with respx.mock:
|
|
respx.post(f"{streamer.base_url}/v2/connections/billing/anycost/test-connection/billing_drops").respond(
|
|
status_code=400, text=marker
|
|
)
|
|
with pytest.raises(httpx.HTTPStatusError) as captured:
|
|
streamer._send_daily_batch("2026-01-01", pl.DataFrame({"cost/cost": [1.0]}), "replace_hourly")
|
|
assert captured.value.response.text == marker
|
|
assert (marker in capsys.readouterr().out) is not redact
|
|
|
|
|
|
@pytest.mark.parametrize("redact", (False, True))
|
|
def test_cloudzero_skips_invalid_timestamp_without_logging_content_when_private(
|
|
redact: bool, capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", redact)
|
|
marker: Final = "synthetic-private-invalid-timestamp"
|
|
streamer: Final = CloudZeroStreamer("test-key", "test-connection")
|
|
data: Final = pl.DataFrame(
|
|
{"time/usage_start": [marker, "2026-01-01T00:00:00Z"], "cost/cost": [9.0, 1.0]}
|
|
)
|
|
with respx.mock:
|
|
route: Final = respx.post(
|
|
f"{streamer.base_url}/v2/connections/billing/anycost/test-connection/billing_drops"
|
|
).respond(status_code=200)
|
|
streamer.send_batched(data)
|
|
assert len(route.calls) == 1
|
|
assert json.loads(route.calls[0].request.content) == {
|
|
"month": "2026-01",
|
|
"operation": "replace_hourly",
|
|
"data": [{"time/usage_start": "2026-01-01T00:00:00Z", "cost/cost": "1"}],
|
|
}
|
|
assert (marker in capsys.readouterr().out) is not redact
|
|
|
|
|
|
@pytest.mark.parametrize("redact", (False, True))
|
|
def test_cloudzero_network_error_preserved_without_logging_content_when_private(
|
|
redact: bool, capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", redact)
|
|
marker: Final = "synthetic-private-network-error"
|
|
streamer: Final = CloudZeroStreamer("test-key", "test-connection")
|
|
failure: Final = httpx.ConnectError(marker)
|
|
with respx.mock:
|
|
respx.post(
|
|
f"{streamer.base_url}/v2/connections/billing/anycost/test-connection/billing_drops"
|
|
).mock(side_effect=failure)
|
|
with pytest.raises(httpx.ConnectError) as captured:
|
|
streamer.send_batched(
|
|
pl.DataFrame({"time/usage_start": ["2026-01-01T00:00:00Z"], "cost/cost": [1.0]})
|
|
)
|
|
assert captured.value is failure
|
|
assert (marker in capsys.readouterr().out) is not redact
|