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>
This commit is contained in:
kerry 2026-09-19 21:24:56 +00:00
parent c0c5cc84f8
commit 5b97d98b7b
4 changed files with 102 additions and 79 deletions

View file

@ -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)

View file

@ -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):

View file

@ -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
}
}
]

View file

@ -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)