from __future__ import annotations from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock import pytest from src.services.usage._billing_integration import UsageBillingIntegrationMixin from src.services.usage._types import UsageRecordParams class _TestUsageBillingIntegration(UsageBillingIntegrationMixin): @classmethod async def _get_rate_multiplier_and_free_tier( cls, db: Any, # noqa: ARG003 provider_api_key_id: str | None, # noqa: ARG003 provider_id: str | None, # noqa: ARG003 api_format: str | None = None, # noqa: ARG003 ) -> tuple[float, bool]: return 1.0, False class _DummyBillingService: last_dimensions: dict[str, Any] | None = None def __init__(self, db: Any) -> None: # noqa: D107, ARG002 pass def calculate( self, *, task_type: str, # noqa: ARG002 model: str, # noqa: ARG002 provider_id: str, # noqa: ARG002 dimensions: dict[str, Any], strict_mode: bool | None, # noqa: ARG002 ) -> Any: _DummyBillingService.last_dimensions = dict(dimensions) snapshot = SimpleNamespace( cost_breakdown={ "input_cost": 0.0, "output_cost": 0.0, "cache_creation_cost": 0.0, "cache_read_cost": 0.0, "request_cost": 0.0, }, total_cost=0.0, resolved_variables={}, to_dict=lambda: {}, ) return SimpleNamespace(snapshot=snapshot) def _build_params( db: Any, *, cache_creation_input_tokens: int = 0, cache_read_input_tokens: int = 0, cache_creation_input_tokens_5m: int = 0, cache_creation_input_tokens_1h: int = 0, cache_ttl_minutes: int | None = None, provider_api_key_id: str | None = "pak-test", ) -> UsageRecordParams: return UsageRecordParams( db=db, user=None, api_key=None, provider="provider-x", model="claude-sonnet", input_tokens=100, output_tokens=50, cache_creation_input_tokens=cache_creation_input_tokens, cache_read_input_tokens=cache_read_input_tokens, cache_creation_input_tokens_5m=cache_creation_input_tokens_5m, cache_creation_input_tokens_1h=cache_creation_input_tokens_1h, request_type="chat", api_format="claude:chat", api_family="claude", endpoint_kind="chat", endpoint_api_format="claude:chat", has_format_conversion=False, is_stream=False, response_time_ms=123, first_byte_time_ms=None, status_code=200, error_message=None, metadata={}, request_headers=None, request_body=None, provider_request_headers=None, provider_request_body=None, response_headers=None, client_response_headers=None, response_body=None, client_response_body=None, request_id="req-test", provider_id="provider-id", provider_endpoint_id="endpoint-id", provider_api_key_id=provider_api_key_id, status="completed", cache_ttl_minutes=cache_ttl_minutes, use_tiered_pricing=True, target_model=None, ) @pytest.mark.asyncio async def test_prepare_usage_record_uses_provider_key_ttl_for_cache_read( monkeypatch: pytest.MonkeyPatch, ) -> None: db = MagicMock() db.query.return_value.filter.return_value.scalar.return_value = 60 monkeypatch.setattr("src.services.billing.service.BillingService", _DummyBillingService) monkeypatch.setattr( "src.services.usage._billing_integration.sanitize_request_metadata", lambda metadata: metadata, ) monkeypatch.setattr( "src.services.usage._billing_integration.build_usage_params", lambda **kwargs: {"total_cost_usd": 0.0, "actual_total_cost_usd": 0.0}, ) params = _build_params(db, cache_read_input_tokens=321) await _TestUsageBillingIntegration._prepare_usage_record(params) assert _DummyBillingService.last_dimensions is not None assert _DummyBillingService.last_dimensions.get("cache_ttl_minutes") == 60 @pytest.mark.asyncio async def test_prepare_usage_record_prefers_explicit_cache_ttl( monkeypatch: pytest.MonkeyPatch, ) -> None: db = MagicMock() db.query.return_value.filter.return_value.scalar.return_value = 60 monkeypatch.setattr("src.services.billing.service.BillingService", _DummyBillingService) monkeypatch.setattr( "src.services.usage._billing_integration.sanitize_request_metadata", lambda metadata: metadata, ) monkeypatch.setattr( "src.services.usage._billing_integration.build_usage_params", lambda **kwargs: {"total_cost_usd": 0.0, "actual_total_cost_usd": 0.0}, ) params = _build_params(db, cache_read_input_tokens=123, cache_ttl_minutes=5) await _TestUsageBillingIntegration._prepare_usage_record(params) assert _DummyBillingService.last_dimensions is not None assert _DummyBillingService.last_dimensions.get("cache_ttl_minutes") == 5 @pytest.mark.asyncio async def test_prepare_usage_record_infers_ttl_from_1h_cache_split( monkeypatch: pytest.MonkeyPatch, ) -> None: db = MagicMock() monkeypatch.setattr("src.services.billing.service.BillingService", _DummyBillingService) monkeypatch.setattr( "src.services.usage._billing_integration.sanitize_request_metadata", lambda metadata: metadata, ) monkeypatch.setattr( "src.services.usage._billing_integration.build_usage_params", lambda **kwargs: {"total_cost_usd": 0.0, "actual_total_cost_usd": 0.0}, ) params = _build_params( db, provider_api_key_id=None, cache_creation_input_tokens=1000, cache_creation_input_tokens_1h=1000, ) await _TestUsageBillingIntegration._prepare_usage_record(params) assert _DummyBillingService.last_dimensions is not None assert _DummyBillingService.last_dimensions.get("cache_ttl_minutes") == 60 @pytest.mark.asyncio async def test_prepare_usage_record_deserializes_body_json_before_build_params( monkeypatch: pytest.MonkeyPatch, ) -> None: db = MagicMock() monkeypatch.setattr("src.services.billing.service.BillingService", _DummyBillingService) monkeypatch.setattr( "src.services.usage._billing_integration.sanitize_request_metadata", lambda metadata: metadata, ) captured: dict[str, Any] = {} def _capture_build_usage_params(**kwargs: Any) -> dict[str, Any]: captured.update(kwargs) return {"total_cost_usd": 0.0, "actual_total_cost_usd": 0.0} monkeypatch.setattr( "src.services.usage._billing_integration.build_usage_params", _capture_build_usage_params, ) params = _build_params(db) params.request_body = '{"messages":[{"role":"user","content":"hello"}]}' params.provider_request_body = '{"tools":[{"name":"calc"}]}' params.response_body = '{"choices":[{"index":0}]}' params.client_response_body = '{"output":[{"type":"text"}]}' await _TestUsageBillingIntegration._prepare_usage_record(params) assert isinstance(captured["request_body"], dict) assert captured["request_body"]["messages"][0]["content"] == "hello" assert isinstance(captured["provider_request_body"], dict) assert isinstance(captured["response_body"], dict) assert isinstance(captured["client_response_body"], dict) @pytest.mark.asyncio async def test_prepare_usage_record_keeps_invalid_json_body_string( monkeypatch: pytest.MonkeyPatch, ) -> None: db = MagicMock() monkeypatch.setattr("src.services.billing.service.BillingService", _DummyBillingService) monkeypatch.setattr( "src.services.usage._billing_integration.sanitize_request_metadata", lambda metadata: metadata, ) captured: dict[str, Any] = {} def _capture_build_usage_params(**kwargs: Any) -> dict[str, Any]: captured.update(kwargs) return {"total_cost_usd": 0.0, "actual_total_cost_usd": 0.0} monkeypatch.setattr( "src.services.usage._billing_integration.build_usage_params", _capture_build_usage_params, ) invalid_json = '{"content":"x...[truncated]' params = _build_params(db) params.request_body = invalid_json await _TestUsageBillingIntegration._prepare_usage_record(params) assert captured["request_body"] == invalid_json