litellm/tests/unit/integrations/focus/test_export_engine.py
devin-ai-integration[bot] c3a23fe499
refactor: expose core private helpers under public names (#44871)
* refactor: expose core private symbols with compatibility aliases

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* types: narrow core migration diagnostics

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix: preserve runtime behavior in core symbol migration

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix: preserve core private usage migration behavior

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix: preserve private value rebinding compatibility

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix: restore optional imports and cover public helpers

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test: exempt router property from call coverage

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test: align recursive detector ignore names

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: mateo <mateo@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-10-07 00:05:27 +00:00

70 lines
2.7 KiB
Python

from datetime import datetime, timezone
import polars as pl
from litellm.integrations.focus.database import FocusLiteLLMDatabase
from litellm.integrations.focus.destinations import FocusTimeWindow
from litellm.integrations.focus.destinations.factory import FocusDestinationFactory
from litellm.integrations.focus.export_engine import FocusExportEngine
from litellm.integrations.focus.serializers import FocusParquetSerializer
from litellm.integrations.focus.transformer import FocusTransformer
def _vantage_csv_engine() -> FocusExportEngine:
return FocusExportEngine(
provider="vantage",
export_format="csv",
prefix="exports",
destination_config={"api_key": "test-key", "integration_token": "test-token"},
)
def test_legacy_engine_properties_forward_to_public_storage():
engine = _vantage_csv_engine()
replacement_destination = FocusDestinationFactory.create(
provider="vantage",
prefix="replacement",
config={"api_key": "test-key", "integration_token": "test-token"},
)
engine._destination = replacement_destination
assert engine._destination is engine.destination is replacement_destination
replacement_serializer = FocusParquetSerializer()
engine._serializer = replacement_serializer
assert engine._serializer is engine.serializer is replacement_serializer
replacement_transformer = FocusTransformer()
engine._transformer = replacement_transformer
assert engine._transformer is engine.transformer is replacement_transformer
replacement_database = FocusLiteLLMDatabase()
engine._database = replacement_database
assert engine._database is engine.database is replacement_database
def test_build_filename_uses_the_window_and_serializer_format():
engine = _vantage_csv_engine()
window = FocusTimeWindow(
start_time=datetime(2024, 1, 2, 3, 4, 5, tzinfo=timezone.utc),
end_time=datetime(2024, 1, 2, 4, 5, 6, tzinfo=timezone.utc),
frequency="hourly",
)
assert engine.build_filename(window) == "usage_20240102T030405Z_20240102T040506Z.csv"
def test_export_aggregates_compute_values_and_handle_missing_columns():
frame = pl.DataFrame(
{
"spend": [1.5, 2.5, 1.0],
"team_id": ["team-a", "team-a", "team-b"],
}
)
assert FocusExportEngine.sum_column(frame, "spend") == 5.0
assert FocusExportEngine.count_unique(frame, "team_id") == 2
assert FocusExportEngine.sum_column(frame, "missing") == 0.0
assert FocusExportEngine.count_unique(frame, "missing") == 0
assert FocusExportEngine.sum_column(pl.DataFrame(), "spend") == 0.0
assert FocusExportEngine.count_unique(pl.DataFrame(), "team_id") == 0