From 5b97d98b7b43136e8192786e5a3fd910ed3229be Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 21:24:56 +0000 Subject: [PATCH] test(integration): require persisted cost breakdowns unless a case opts out Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integration/cost_calculation/conftest.py | 11 +- .../cost_calculation/cost_tracking_case.py | 1 + .../cost_calculation/cost_tracking_cases.json | 24 ++- .../cost_calculation/test_cost_tracking.py | 145 ++++++++++-------- 4 files changed, 102 insertions(+), 79 deletions(-) diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index 173d48b052d..4d4c7b6356a 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -40,6 +40,7 @@ class CostRow(BaseModel): model_config = ConfigDict(extra="ignore") spend: float | None = None + status: str | None = None prompt_tokens: int | None = None completion_tokens: int | None = None metadata: CostMetadata | None = None @@ -62,10 +63,7 @@ def approx_equal(actual: float, expected: float) -> bool: return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) -def assert_total_is_sum_of_components(row: CostRow, context: str) -> None: - breakdown: Final = row.breakdown - if breakdown is None: - return +def assert_total_is_sum_of_components(row: CostRow, breakdown: CostBreakdown, context: str) -> None: total: Final = sum( cost or 0.0 for cost in (breakdown.input_cost, breakdown.output_cost, breakdown.tool_usage_cost) @@ -84,7 +82,7 @@ def _row(value: Mapping[str, object]) -> CostRow | None: metadata_value: Final = value.get("metadata") metadata: Final = json.loads(metadata_value) if isinstance(metadata_value, str) else metadata_value parsed: Final = CostRow.model_validate({**value, "metadata": metadata}) - return parsed + return parsed if parsed.metadata is not None or (parsed.spend is not None and parsed.status is not None) else None def poll_cost_row(key: str) -> CostRow: @@ -92,7 +90,8 @@ def poll_cost_row(key: str) -> CostRow: def read() -> CostRow | None: rows: Final = read_rows( - 'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + 'SELECT spend, status, metadata, prompt_tokens, completion_tokens ' + 'FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,), ) return next((parsed for row in rows if (parsed := _row(row)) is not None), None) diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index 23180ad16ef..376477f44f1 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -143,6 +143,7 @@ class ExactExpected(BaseModel): cache_creation_cost: float | None = None reasoning_cost: float | None = None tool_usage_cost: float | None = None + breakdown_persisted: bool = True class RecountRates(BaseModel): diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index 40f9c7a6762..b7da4642cc8 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -26120,7 +26120,8 @@ "input_cost": 0.04, "output_cost": 0, "prompt_tokens": 0, - "completion_tokens": 0 + "completion_tokens": 0, + "breakdown_persisted": false } }, { @@ -26153,7 +26154,8 @@ "input_cost": 0.08, "output_cost": 0, "prompt_tokens": 0, - "completion_tokens": 0 + "completion_tokens": 0, + "breakdown_persisted": false } }, { @@ -26186,7 +26188,8 @@ "input_cost": 0.06, "output_cost": 0, "prompt_tokens": 0, - "completion_tokens": 0 + "completion_tokens": 0, + "breakdown_persisted": false } }, { @@ -26222,7 +26225,8 @@ "input_cost": 0.08, "output_cost": 0, "prompt_tokens": 0, - "completion_tokens": 0 + "completion_tokens": 0, + "breakdown_persisted": false } }, { @@ -26264,7 +26268,8 @@ "input_cost": 1.71e-05, "output_cost": 0.000102, "prompt_tokens": 10, - "completion_tokens": 20 + "completion_tokens": 20, + "breakdown_persisted": false } }, { @@ -26292,7 +26297,8 @@ "input_cost": 0.05, "output_cost": 0, "prompt_tokens": 0, - "completion_tokens": 0 + "completion_tokens": 0, + "breakdown_persisted": false } }, { @@ -26319,7 +26325,8 @@ "input_cost": 0.045, "output_cost": 0, "prompt_tokens": 0, - "completion_tokens": 0 + "completion_tokens": 0, + "breakdown_persisted": false } }, { @@ -26364,7 +26371,8 @@ "input_cost": 1.7e-05, "output_cost": 0.000102, "prompt_tokens": 10, - "completion_tokens": 20 + "completion_tokens": 20, + "breakdown_persisted": false } } ] diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index 944d903625b..3e314c66b74 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -15,6 +15,7 @@ import pytest from integration._support.client import JSON_OBJECT, Gateway from integration.cost_calculation.conftest import ( + CostBreakdown, approx_equal, assert_total_is_sum_of_components, poll_cost_row, @@ -93,6 +94,76 @@ def _assert_stream_has_no_error(response_text: str) -> None: assert "error" not in parsed, f"stream carried an error event: {parsed}" +def _assert_breakdown( + case: CostTrackingTestCase, + expected: ExactExpected, + breakdown: CostBreakdown, + response: httpx.Response, +) -> None: + assert breakdown.input_cost is not None and approx_equal(breakdown.input_cost, expected.input_cost), ( + f"{case.name}: input_cost {breakdown.input_cost} != expected {expected.input_cost}" + ) + assert breakdown.output_cost is not None and approx_equal(breakdown.output_cost, expected.output_cost), ( + f"{case.name}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" + ) + for field, header_name, actual_component, expected_component in ( + ( + "cache_read_cost", + "x-litellm-response-cost-cache-read", + breakdown.cache_read_cost, + expected.cache_read_cost, + ), + ( + "cache_creation_cost", + "x-litellm-response-cost-cache-creation", + breakdown.cache_creation_cost, + expected.cache_creation_cost, + ), + ( + "reasoning_cost", + "x-litellm-response-cost-reasoning", + breakdown.reasoning_cost, + expected.reasoning_cost, + ), + ( + "tool_usage_cost", + "x-litellm-response-cost-tool-usage", + breakdown.tool_usage_cost, + expected.tool_usage_cost, + ), + ): + if expected_component is None: + continue + assert actual_component is not None and approx_equal(actual_component, expected_component), ( + f"{case.name}: {field} {actual_component} != expected {expected_component}" + ) + if case.response.content_type == "application/json": + header: Final = response.headers.get(header_name) + assert header is not None and approx_equal(float(header), expected_component), ( + f"{case.name}: {header_name} {header} != expected {expected_component}" + ) + if case.response.content_type == "application/json" and any( + component is not None + for component in ( + expected.cache_read_cost, + expected.cache_creation_cost, + expected.reasoning_cost, + expected.tool_usage_cost, + ) + ): + input_header: Final = response.headers.get("x-litellm-response-cost-input") + output_header: Final = response.headers.get("x-litellm-response-cost-output") + expected_input_header: Final = expected.input_cost - ( + expected.cache_read_cost or 0.0 + ) - (expected.cache_creation_cost or 0.0) + assert input_header is not None and approx_equal(float(input_header), expected_input_header), ( + f"{case.name}: x-litellm-response-cost-input {input_header} != expected {expected_input_header}" + ) + assert output_header is not None and approx_equal(float(output_header), expected.output_cost), ( + f"{case.name}: x-litellm-response-cost-output {output_header} != expected {expected.output_cost}" + ) + + @pytest.mark.parametrize("case", _CASES) def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) -> None: marker: Final = sha256(case.name.encode()).hexdigest()[:12] @@ -133,7 +204,9 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert row.spend is not None and approx_equal(row.spend, recount), ( f"{case.name}: spend {row.spend} != recount {recount} at map rates" ) - assert_total_is_sum_of_components(row, case.name) + breakdown: Final = row.breakdown + assert breakdown is not None, f"{case.name}: no cost_breakdown persisted" + assert_total_is_sum_of_components(row, breakdown, case.name) return expected: Final = case.expected assert isinstance(expected, ExactExpected) @@ -150,76 +223,18 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) ) assert row.spend is not None and approx_equal(row.spend, expected.spend), ( f"{case.name}: spend {row.spend} != expected {expected.spend} " - f"(breakdown {row.breakdown.model_dump()})" + f"(breakdown {row.breakdown.model_dump() if row.breakdown is not None else None})" ) breakdown: Final = row.breakdown + if expected.breakdown_persisted: + assert breakdown is not None, f"{case.name}: no cost_breakdown persisted" if breakdown is not None: - assert breakdown.input_cost is not None and approx_equal(breakdown.input_cost, expected.input_cost), ( - f"{case.name}: input_cost {breakdown.input_cost} != expected {expected.input_cost}" - ) - assert breakdown.output_cost is not None and approx_equal(breakdown.output_cost, expected.output_cost), ( - f"{case.name}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" - ) - for field, header_name, actual_component, expected_component in ( - ( - "cache_read_cost", - "x-litellm-response-cost-cache-read", - breakdown.cache_read_cost, - expected.cache_read_cost, - ), - ( - "cache_creation_cost", - "x-litellm-response-cost-cache-creation", - breakdown.cache_creation_cost, - expected.cache_creation_cost, - ), - ( - "reasoning_cost", - "x-litellm-response-cost-reasoning", - breakdown.reasoning_cost, - expected.reasoning_cost, - ), - ( - "tool_usage_cost", - "x-litellm-response-cost-tool-usage", - breakdown.tool_usage_cost, - expected.tool_usage_cost, - ), - ): - if expected_component is None: - continue - assert actual_component is not None and approx_equal(actual_component, expected_component), ( - f"{case.name}: {field} {actual_component} != expected {expected_component}" - ) - if case.response.content_type == "application/json": - header: Final = response.headers.get(header_name) - assert header is not None and approx_equal(float(header), expected_component), ( - f"{case.name}: {header_name} {header} != expected {expected_component}" - ) - if case.response.content_type == "application/json" and any( - component is not None - for component in ( - expected.cache_read_cost, - expected.cache_creation_cost, - expected.reasoning_cost, - expected.tool_usage_cost, - ) - ): - input_header: Final = response.headers.get("x-litellm-response-cost-input") - output_header: Final = response.headers.get("x-litellm-response-cost-output") - expected_input_header: Final = expected.input_cost - ( - expected.cache_read_cost or 0.0 - ) - (expected.cache_creation_cost or 0.0) - assert input_header is not None and approx_equal(float(input_header), expected_input_header), ( - f"{case.name}: x-litellm-response-cost-input {input_header} != expected {expected_input_header}" - ) - assert output_header is not None and approx_equal(float(output_header), expected.output_cost), ( - f"{case.name}: x-litellm-response-cost-output {output_header} != expected {expected.output_cost}" - ) + _assert_breakdown(case, expected, breakdown, response) assert row.prompt_tokens == expected.prompt_tokens, ( f"{case.name}: prompt_tokens {row.prompt_tokens} != expected {expected.prompt_tokens}" ) assert row.completion_tokens == expected.completion_tokens, ( f"{case.name}: completion_tokens {row.completion_tokens} != expected {expected.completion_tokens}" ) - assert_total_is_sum_of_components(row, case.name) + if breakdown is not None: + assert_total_is_sum_of_components(row, breakdown, case.name)