From a841750d460fdd148aa761c62c494d989c744ea9 Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 00:19:56 +0000 Subject: [PATCH 1/7] test(integration): proxy behaviour cost cases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/_support/upstream.py | 12 +- tests/integration/contracts.json | 42 ++ .../integration/cost_calculation/conftest.py | 113 +++- .../cost_calculation/cost_tracking_case.py | 34 +- .../cost_calculation/cost_tracking_cases.json | 624 +++++++++++++++++- .../cost_calculation/test_cost_tracking.py | 191 +++++- 6 files changed, 965 insertions(+), 51 deletions(-) diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index 0bea824e77b..acaf036d507 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -1,9 +1,10 @@ from __future__ import annotations import argparse +import asyncio import base64 from collections import deque -from collections.abc import Mapping +from collections.abc import AsyncIterator, Mapping import json from dataclasses import dataclass, field import os @@ -18,7 +19,7 @@ import uvicorn from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from starlette.applications import Starlette from starlette.requests import Request -from starlette.responses import JSONResponse, Response +from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations @@ -223,6 +224,13 @@ class Provider: media_type=response.content_type, ) case SseResponse(): + if response.frame_delay_ms > 0: + async def stream() -> AsyncIterator[bytes]: + for frame in response.frames: + yield f"{frame.replace('$REQUEST_ID', scenario_id)}\n\n".encode() + await asyncio.sleep(response.frame_delay_ms / 1000) + + return StreamingResponse(stream(), media_type=response.content_type) stream_body: Final = ("\n\n".join(response.frames) + "\n\n").replace( "$REQUEST_ID", scenario_id ) diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index edcc7124f38..fe7c808790d 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -1539,6 +1539,48 @@ ], "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[openai-deployment-pricing-override]": [ "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-upstream_400_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-upstream_401_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-upstream_500_stream_request_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-responses_upstream_500_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[claude-sonnet-5-messages_upstream_500_zero_spend]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-fallback_billed_to_answering_deployment]": [ + "quota_management.spend_tracking.routing.fallback_billing" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-n_2_choices]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-finish_reason_length]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-stream_usage_in_empty_choices_chunk]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-stream_usage_in_last_delta_chunk]": [ + "quota_management.spend_tracking.scripted_wire.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-unknown_model_response_model_unknown]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-unknown_model_response_model_known]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-chat_request_to_embedding_entry]": [ + "quota_management.spend_tracking.cost_matrix.logs_cost" + ], + "tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[gpt-5.6-client_disconnect_mid_stream]": [ + "quota_management.spend_tracking.scripted_wire.client_disconnect" ] }, "browser": { diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index b9dae412fa1..a70cd3619ae 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -4,6 +4,7 @@ import functools import json import os from collections.abc import Mapping +from dataclasses import dataclass from hashlib import sha256 from typing import Final @@ -13,8 +14,8 @@ from pydantic import BaseModel, ConfigDict from integration._support.client import JSON_OBJECT, Scenario, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.upstream import delete_scenario, register_scenario -from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase +from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase, StoredResponse class CostBreakdown(BaseModel): @@ -43,6 +44,7 @@ class CostRow(BaseModel): status: str | None = None prompt_tokens: int | None = None completion_tokens: int | None = None + model_id: str | None = None metadata: CostMetadata | None = None @property @@ -59,6 +61,26 @@ class FailureRow(BaseModel): completion_tokens: int | None = None +class DailySpend(BaseModel): + model_config = ConfigDict(extra="ignore") + + spend: float | None = None + prompt_tokens: int | None = None + completion_tokens: int | None = None + api_requests: int | None = None + + +class Rollups(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + key_spend: float + team_spend: float + user_spend: float + end_user_spend: float + daily_user: DailySpend + daily_team: DailySpend + + def approx_equal(actual: float, expected: float) -> bool: return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) @@ -90,7 +112,7 @@ def poll_cost_row(key: str) -> CostRow: def read() -> CostRow | None: rows: Final = read_rows( - 'SELECT spend, status, metadata, prompt_tokens, completion_tokens ' + 'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id ' 'FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,), ) @@ -101,6 +123,67 @@ def poll_cost_row(key: str) -> CostRow: return result +def poll_rows(key: str) -> tuple[CostRow, ...]: + digest: Final = sha256(key.encode()).hexdigest() + + def read() -> tuple[CostRow, ...]: + rows: Final = read_rows( + 'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id ' + 'FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"', + (digest,), + ) + return tuple(parsed for row in rows if (parsed := _row(row)) is not None) + + result: Final = eventually(read, lambda rows: len(rows) > 0, seconds=20) + return result + + +def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str) -> Rollups: + digest: Final = sha256(key.encode()).hexdigest() + + def read() -> Rollups | None: + key_rows: Final = read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', + (digest,), + ) + team_rows: Final = read_rows( + 'SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id=%s', + (team_id,), + ) + user_rows: Final = read_rows( + 'SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=%s', + (user_id,), + ) + end_user_rows: Final = read_rows( + 'SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', + (end_user_id,), + ) + daily_user_rows: Final = read_rows( + 'SELECT spend, prompt_tokens, completion_tokens, api_requests ' + 'FROM "LiteLLM_DailyUserSpend" WHERE user_id=%s AND api_key=%s AND date=CURRENT_DATE::text', + (user_id, digest), + ) + daily_team_rows: Final = read_rows( + 'SELECT spend, prompt_tokens, completion_tokens, api_requests ' + 'FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s AND api_key=%s AND date=CURRENT_DATE::text', + (team_id, digest), + ) + if not all((key_rows, team_rows, user_rows, end_user_rows, daily_user_rows, daily_team_rows)): + return None + return Rollups( + key_spend=float(key_rows[0]["spend"]), + team_spend=float(team_rows[0]["spend"]), + user_spend=float(user_rows[0]["spend"]), + end_user_spend=float(end_user_rows[0]["spend"]), + daily_user=DailySpend.model_validate(daily_user_rows[0]), + daily_team=DailySpend.model_validate(daily_team_rows[0]), + ) + + result: Final = eventually(read, lambda value: value is not None, seconds=20) + assert result is not None + return result + + def poll_failure_row(key: str) -> FailureRow: digest: Final = sha256(key.encode()).hexdigest() @@ -147,17 +230,31 @@ def _vertex_service_account_json(url: str) -> str: ) +@dataclass(frozen=True, slots=True) +class RegisteredDeployment: + model_name: str + identity: str + handle: ScenarioHandle + + def register_scenario_deployment( scenario: Scenario, case: CostTrackingTestCase, marker: str, key: str, -) -> str: + *, + response: StoredResponse | None = None, + marker_suffix: str = "", + model_name: str | None = None, +) -> RegisteredDeployment: control_url: Final = os.environ["INTEGRATION_UPSTREAM_URL"].rstrip("/") run_marker: Final = sha256(key.encode()).hexdigest()[:12] - handle: Final = register_scenario(f"sc-{marker}-{run_marker}", case.response) + handle: Final = register_scenario( + f"sc-{marker}{marker_suffix}-{run_marker}", + case.response if response is None else response, + ) scenario.cleanups.callback(delete_scenario, handle) - model_name: Final = f"cost-{marker}-{run_marker}" + registered_model_name: Final = model_name or f"cost-{marker}{marker_suffix}-{run_marker}" parameters: Final = { "model": case.litellm_model, "api_key": case.api_key, @@ -184,7 +281,7 @@ def register_scenario_deployment( created: Final = scenario.gateway.post( "/model/new", JSON_OBJECT.validate_python({ - "model_name": model_name, + "model_name": registered_model_name, "litellm_params": parameters, "model_info": ( {"base_model": case.base_model} @@ -195,4 +292,4 @@ def register_scenario_deployment( ) identity: Final = string_value(object_value(created["model_info"])["id"]) scenario.cleanups.callback(scenario.delete_model, identity) - return model_name + return RegisteredDeployment(model_name=registered_model_name, identity=identity, handle=handle) diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index 456148bc3fe..f111778b233 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -116,6 +116,7 @@ class SseResponse(BaseModel): content_type: Literal["text/event-stream"] frames: tuple[str, ...] + frame_delay_ms: int = Field(default=0, ge=0) class EventStreamEvent(BaseModel): @@ -160,6 +161,7 @@ class ExactExpected(BaseModel): tool_usage_cost: float | None = None breakdown_persisted: bool = True cost_header: bool = True + rollups: bool = False class RecountRates(BaseModel): @@ -173,6 +175,8 @@ class RecountExpected(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") recount: RecountRates + prompt_tokens: int | None = None + completion_tokens: int | None = None class FailureDetails(BaseModel): @@ -217,6 +221,8 @@ class CostTrackingTestCase(BaseModel): request: dict[str, JsonValue] response: StoredResponse expected: Expected + fallback_from: StoredResponse | None = None + disconnect_after_frames: int | None = Field(default=None, ge=1) @property def rates(self) -> CostMapEntry: @@ -430,7 +436,31 @@ def data_errors() -> tuple[str, ...]: and case.rates.mode != "image_generation" and not case.reports_provider_cost ) - or (not case.expected.cost_header and case.passthrough_provider is None) + or ( + not case.expected.cost_header + and case.passthrough_provider is None + and not isinstance(case.response, SseResponse) + and case.expected.spend != 0.0 + ) + ) + ) + invalid_fallbacks: Final = sorted( + case.name + for case in CASES + if case.fallback_from is not None + and ( + not isinstance(case.fallback_from, JsonResponse) + or not 400 <= case.fallback_from.status <= 599 + ) + ) + invalid_disconnects: Final = sorted( + case.name + for case in CASES + if case.disconnect_after_frames is not None + and ( + not isinstance(case.response, SseResponse) + or case.response.frame_delay_ms <= 0 + or not isinstance(case.expected, RecountExpected) ) ) return tuple( @@ -446,6 +476,8 @@ def data_errors() -> tuple[str, ...]: if failure_response_mismatches else None, f"invalid passthrough opt-outs: {invalid_opt_outs}" if invalid_opt_outs else None, + f"invalid fallback responses: {invalid_fallbacks}" if invalid_fallbacks else None, + f"invalid disconnect cases: {invalid_disconnects}" if invalid_disconnects else None, ) if message is not None ) diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index d532a6851fd..9fb7c6990a7 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -617,6 +617,11 @@ "mode": "chat", "input_cost_per_token": 1.51e-06, "output_cost_per_token": 7.51e-06 + }, + "text-embedding-3-large": { + "litellm_provider": "openai", + "mode": "embedding", + "input_cost_per_token": 1.3e-07 } }, "cases": [ @@ -6947,7 +6952,8 @@ "input_cost": 0.00552, "output_cost": 0.00618, "prompt_tokens": 1840, - "completion_tokens": 412 + "completion_tokens": 412, + "rollups": true } }, { @@ -7608,7 +7614,9 @@ "recount": { "input_cost_per_token": 3e-06, "output_cost_per_token": 1.5e-05 - } + }, + "prompt_tokens": 47, + "completion_tokens": 10 } }, { @@ -7695,7 +7703,9 @@ "recount": { "input_cost_per_token": 3e-06, "output_cost_per_token": 1.5e-05 - } + }, + "prompt_tokens": 49, + "completion_tokens": 111 } }, { @@ -7752,7 +7762,9 @@ "recount": { "input_cost_per_token": 3e-06, "output_cost_per_token": 1.5e-05 - } + }, + "prompt_tokens": 301, + "completion_tokens": 9 } }, { @@ -12782,7 +12794,9 @@ "recount": { "input_cost_per_token": 5.2e-07, "output_cost_per_token": 3.12e-06 - } + }, + "prompt_tokens": 48, + "completion_tokens": 12 } }, { @@ -12862,7 +12876,9 @@ "recount": { "input_cost_per_token": 5.2e-07, "output_cost_per_token": 3.12e-06 - } + }, + "prompt_tokens": 46, + "completion_tokens": 88 } }, { @@ -12914,7 +12930,9 @@ "recount": { "input_cost_per_token": 5.2e-07, "output_cost_per_token": 3.12e-06 - } + }, + "prompt_tokens": 302, + "completion_tokens": 10 } }, { @@ -20760,7 +20778,8 @@ "input_cost": 0.00322, "output_cost": 0.005768, "prompt_tokens": 1840, - "completion_tokens": 412 + "completion_tokens": 412, + "rollups": true } }, { @@ -21515,7 +21534,9 @@ "recount": { "input_cost_per_token": 1.75e-06, "output_cost_per_token": 1.4e-05 - } + }, + "prompt_tokens": 49, + "completion_tokens": 12 } }, { @@ -21601,7 +21622,9 @@ "recount": { "input_cost_per_token": 1.75e-06, "output_cost_per_token": 1.4e-05 - } + }, + "prompt_tokens": 44, + "completion_tokens": 105 } }, { @@ -21656,7 +21679,9 @@ "recount": { "input_cost_per_token": 1.75e-06, "output_cost_per_token": 1.4e-05 - } + }, + "prompt_tokens": 302, + "completion_tokens": 11 } }, { @@ -29417,6 +29442,583 @@ "prompt_tokens": 1840, "completion_tokens": 412 } + }, + { + "name": "gpt-5.6-upstream_400_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "status": 400, + "body": { + "error": { + "message": "scripted upstream failure 400", + "type": "server_error", + "code": "400" + } + } + }, + "expected": { + "failure": { + "status": 400 + } + } + }, + { + "name": "gpt-5.6-upstream_401_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "status": 401, + "body": { + "error": { + "message": "scripted upstream failure 401", + "type": "server_error", + "code": "401" + } + } + }, + "expected": { + "failure": { + "status": 401 + } + } + }, + { + "name": "gpt-5.6-upstream_500_stream_request_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": true + }, + "response": { + "content_type": "application/json", + "status": 500, + "body": { + "error": { + "message": "scripted upstream failure 500", + "type": "server_error", + "code": "500" + } + } + }, + "expected": { + "failure": { + "status": 500 + } + } + }, + { + "name": "gpt-5.6-responses_upstream_500_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/responses", + "request": { + "model": "$MODEL", + "input": "proxy behaviour probe", + "stream": false + }, + "response": { + "content_type": "application/json", + "status": 500, + "body": { + "error": { + "message": "scripted upstream failure 500", + "type": "server_error", + "code": "500" + } + } + }, + "expected": { + "failure": { + "status": 500 + } + } + }, + { + "name": "claude-sonnet-5-messages_upstream_500_zero_spend", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "endpoint": "/v1/messages", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ] + }, + "response": { + "content_type": "application/json", + "status": 500, + "body": { + "error": { + "message": "scripted upstream failure 500", + "type": "server_error", + "code": "500" + } + } + }, + "expected": { + "failure": { + "status": 500 + } + } + }, + { + "name": "gpt-5.6-fallback_billed_to_answering_deployment", + "covers": "quota_management.spend_tracking.routing.fallback_billing", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1789788262, + "model": "gpt-5.6", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "fallback_from": { + "content_type": "application/json", + "status": 500, + "body": { + "error": { + "message": "scripted upstream failure 500", + "type": "server_error", + "code": "500" + } + } + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-n_2_choices", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1789788262, + "model": "gpt-5.6", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer" + }, + "finish_reason": "stop" + }, + { + "index": 1, + "message": { + "role": "assistant", + "content": "second choice" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-finish_reason_length", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1789788262, + "model": "gpt-5.6", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "truncated" + }, + "finish_reason": "length" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-stream_usage_in_empty_choices_chunk", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": true, + "stream_options": { + "include_usage": true + } + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"scripted answer\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[],\"usage\":{\"prompt_tokens\":1840,\"completion_tokens\":412,\"total_tokens\":2252}}", + "data: [DONE]" + ] + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412, + "cost_header": false + } + }, + { + "name": "gpt-5.6-stream_usage_in_last_delta_chunk", + "covers": "quota_management.spend_tracking.scripted_wire.logs_cost", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": true, + "stream_options": { + "include_usage": true + } + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"scripted answer\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":1840,\"completion_tokens\":412,\"total_tokens\":2252}}", + "data: [DONE]" + ] + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412, + "cost_header": false + } + }, + { + "name": "gpt-5.6-unknown_model_response_model_unknown", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "deployment": { + "model": "openai/not-in-any-map-xyz" + }, + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1789788262, + "model": "not-in-any-map-xyz", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.0, + "input_cost": 0.0, + "output_cost": 0.0, + "prompt_tokens": 1840, + "completion_tokens": 412, + "cost_header": false + } + }, + { + "name": "gpt-5.6-unknown_model_response_model_known", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "gpt-5.6", + "deployment": { + "model": "openai/not-in-any-map-xyz" + }, + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1789788262, + "model": "gpt-5.6", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.008988, + "input_cost": 0.00322, + "output_cost": 0.005768, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-chat_request_to_embedding_entry", + "covers": "quota_management.spend_tracking.cost_matrix.logs_cost", + "model": "text-embedding-3-large", + "deployment": { + "model": "openai/text-embedding-3-large" + }, + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": false + }, + "response": { + "content_type": "application/json", + "body": { + "id": "chatcmpl-$REQUEST_ID", + "object": "chat.completion", + "created": 1789788262, + "model": "text-embedding-3-large", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "scripted answer" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 1840, + "completion_tokens": 412, + "total_tokens": 2252 + } + } + }, + "expected": { + "spend": 0.0002392, + "input_cost": 0.0002392, + "output_cost": 0.0, + "prompt_tokens": 1840, + "completion_tokens": 412 + } + }, + { + "name": "gpt-5.6-client_disconnect_mid_stream", + "covers": "quota_management.spend_tracking.scripted_wire.client_disconnect", + "model": "gpt-5.6", + "request": { + "model": "$MODEL", + "messages": [ + { + "role": "user", + "content": "proxy behaviour probe" + } + ], + "stream": true, + "stream_options": { + "include_usage": true + } + }, + "response": { + "content_type": "text/event-stream", + "frames": [ + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-0\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-1\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-2\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-3\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-4\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-5\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-6\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-7\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-8\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-9\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-10\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-11\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-12\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-13\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-14\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-15\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-16\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-17\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-18\"},\"finish_reason\":null}],\"usage\":null}", + "data: {\"id\":\"chatcmpl-$REQUEST_ID\",\"object\":\"chat.completion.chunk\",\"created\":1789788263,\"model\":\"gpt-5.6\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"frame-19\"},\"finish_reason\":null}],\"usage\":null}", + "data: [DONE]" + ], + "frame_delay_ms": 200 + }, + "disconnect_after_frames": 3, + "expected": { + "recount": { + "input_cost_per_token": 1.75e-06, + "output_cost_per_token": 1.4e-05 + } + } } ] } diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index 8c5f306dd9d..e8eed98205e 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -3,10 +3,12 @@ from __future__ import annotations import io +from itertools import islice import json from hashlib import sha256 import struct from typing import Final, cast +import uuid import wave import zlib @@ -18,10 +20,13 @@ from integration._support.client import JSON_OBJECT, Gateway from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.conftest import ( CostBreakdown, + CostRow, approx_equal, assert_total_is_sum_of_components, poll_cost_row, poll_failure_row, + poll_rollups, + poll_rows, register_scenario_deployment, ) from integration.cost_calculation.cost_tracking_case import ( @@ -179,11 +184,47 @@ def _assert_breakdown( ) +def _assert_exact( + case: CostTrackingTestCase, + expected: ExactExpected, + row: CostRow, + response: httpx.Response, +) -> None: + 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() 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(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}" + ) + if breakdown is not None: + assert_total_is_sum_of_components(row, breakdown, case.name) + + @pytest.mark.parametrize("case", _CASES) def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) -> None: marker: Final = sha256(case.name.encode()).hexdigest()[:12] with gateway.scenario() as scenario: - key: Final = scenario.key() + expected: Final = case.expected + team_id: Final = scenario.team() if isinstance(expected, ExactExpected) and expected.rollups else None + user_id: Final = ( + scenario.user(team_id=team_id) + if team_id is not None + else None + ) + key: Final = ( + scenario.key(team_id=team_id, user_id=user_id) + if team_id is not None and user_id is not None + else scenario.key() + ) passthrough_provider: Final = case.passthrough_provider scenario_id: Final = f"sc-{marker}-{sha256(key.encode()).hexdigest()[:12]}" scenario_handle: Final = ( @@ -193,21 +234,73 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) ) if scenario_handle is not None: scenario.cleanups.callback(delete_scenario, scenario_handle) + deployment: Final = ( + register_scenario_deployment(scenario, case, marker, key) + if passthrough_provider is None + else None + ) + fallback_deployment: Final = ( + register_scenario_deployment( + scenario, + case, + marker, + key, + response=case.fallback_from, + marker_suffix="-fb", + ) + if case.fallback_from is not None + else None + ) + if isinstance(expected, ExactExpected) and expected.rollups: + assert deployment is not None + rollup_deployments: Final = tuple( + register_scenario_deployment( + scenario, + case, + marker, + key, + marker_suffix=f"-r{index}", + model_name=deployment.model_name, + ) + for index in (2, 3) + ) + assert len(rollup_deployments) == 2 model_name: Final = ( case.model if passthrough_provider in {"gemini", "anthropic"} - else register_scenario_deployment(scenario, case, marker, key) + else deployment.model_name if deployment is not None else None ) + assert model_name is not None request_model: Final = ( case.model.rsplit("/", 1)[-1] if passthrough_provider in {"gemini", "anthropic"} - else model_name + else fallback_deployment.model_name if fallback_deployment is not None else model_name ) - request_body: Final = JSON_OBJECT.validate_python( + base_request_values: Final = ( _replace_model(case.request, request_model) if passthrough_provider is not None else {**case.request, "model": model_name} ) + end_user_id: Final = ( + f"end-user-{uuid.uuid4()}" + if isinstance(expected, ExactExpected) and expected.rollups + else None + ) + request_body: Final = JSON_OBJECT.validate_python( + { + **base_request_values, + **( + {"model": fallback_deployment.model_name, "fallbacks": [model_name]} + if fallback_deployment is not None + else {} + ), + **( + {"user": end_user_id, "cache": {"no-cache": True}} + if end_user_id is not None + else {} + ), + } + ) request_headers: Final = ( { "x-pass-x-scripted-scenario": scenario_id, @@ -225,12 +318,45 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) if passthrough_provider is not None else case.endpoint ) - response: Final = ( - _multipart_request(gateway, case, model_name, key) - if case.upload is not None - else gateway.request("POST", request_path, request_body, key=key, headers=request_headers) + if case.disconnect_after_frames is not None: + with gateway.client.stream( + "POST", + request_path, + json=request_body, + headers={"Authorization": f"Bearer {key}", **request_headers}, + ) as stream_response: + frames: Final = tuple( + islice( + (line for line in stream_response.iter_lines() if line.startswith("data:")), + case.disconnect_after_frames, + ) + ) + assert len(frames) == case.disconnect_after_frames + row: Final = poll_cost_row(key) + assert isinstance(expected, RecountExpected) + if expected.prompt_tokens is not None: + assert row.prompt_tokens == expected.prompt_tokens + if expected.completion_tokens is not None: + assert row.completion_tokens == expected.completion_tokens + assert row.prompt_tokens is not None and row.prompt_tokens > 0 + assert row.completion_tokens is not None and row.completion_tokens > 0 + recount: Final = row.prompt_tokens * expected.recount.input_cost_per_token + ( + row.completion_tokens * expected.recount.output_cost_per_token + ) + assert row.spend is not None and approx_equal(row.spend, recount) + assert row.breakdown is not None + assert_total_is_sum_of_components(row, row.breakdown, case.name) + return + responses: Final = tuple( + ( + _multipart_request(gateway, case, model_name, key) + if case.upload is not None + else gateway.request("POST", request_path, request_body, key=key, headers=request_headers) + ) + for _ in range(3 if isinstance(expected, ExactExpected) and expected.rollups else 1) ) - if isinstance(case.expected, FailureExpected): + response: Final = responses[0] + if isinstance(expected, FailureExpected): assert response.status_code == case.expected.failure.status, ( f"{case.name}: proxy returned {response.status_code}, expected {case.expected.failure.status}: " f"{response.text[:400]}" @@ -245,8 +371,9 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert response.is_success, f"{case.name}: proxy returned {response.status_code}: {response.text[:400]}" if case.response.content_type == "text/event-stream": _assert_stream_has_no_error(response.text) - row: Final = poll_cost_row(key) - if isinstance(case.expected, RecountExpected): + rows: Final = poll_rows(key) if len(responses) > 1 else (poll_cost_row(key),) + if isinstance(expected, RecountExpected): + row: Final = rows[0] assert row.prompt_tokens is not None and row.prompt_tokens > 0, ( f"{case.name}: recount case counted no input tokens: prompt_tokens={row.prompt_tokens}" ) @@ -263,8 +390,12 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) 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) + if fallback_deployment is not None: + assert deployment is not None + assert len(rows) == 1 + assert rows[0].status == "success" + assert rows[0].model_id == deployment.identity if isinstance(case.response, BinaryResponse): header: Final = response.headers.get("x-litellm-response-cost") if header is not None: @@ -281,20 +412,22 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert approx_equal(float(header), expected.spend), ( f"{case.name}: x-litellm-response-cost {header} != expected {expected.spend}" ) - 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() 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(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}" - ) - if breakdown is not None: - assert_total_is_sum_of_components(row, breakdown, case.name) + for row in rows: + _assert_exact(case, expected, row, response) + if expected.rollups: + assert deployment is not None and team_id is not None and user_id is not None + assert end_user_id is not None + rollups: Final = poll_rollups(key, team_id, user_id, end_user_id) + target_spend: Final = expected.spend * 3 + assert approx_equal(rollups.key_spend, target_spend) + assert approx_equal(rollups.team_spend, target_spend) + assert approx_equal(rollups.user_spend, target_spend) + assert approx_equal(rollups.end_user_spend, target_spend) + assert approx_equal(rollups.daily_user.spend or 0.0, target_spend) + assert approx_equal(rollups.daily_team.spend or 0.0, target_spend) + assert rollups.daily_user.prompt_tokens == expected.prompt_tokens * 3 + assert rollups.daily_user.completion_tokens == expected.completion_tokens * 3 + assert rollups.daily_user.api_requests == 3 + assert rollups.daily_team.prompt_tokens == expected.prompt_tokens * 3 + assert rollups.daily_team.completion_tokens == expected.completion_tokens * 3 + assert rollups.daily_team.api_requests == 3 From 2edea0be086ebbf9c7c6ae0c1a539b64c588fabe Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 00:36:18 +0000 Subject: [PATCH 2/7] test(integration): assert recount pins and unique fixture request ids Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/_support/upstream.py | 22 ++++-- .../integration/cost_calculation/conftest.py | 13 ++-- .../cost_calculation/cost_tracking_case.py | 22 ++++++ .../cost_calculation/cost_tracking_cases.json | 17 ++-- .../cost_calculation/test_cost_tracking.py | 77 ++++++++----------- 5 files changed, 86 insertions(+), 65 deletions(-) diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index acaf036d507..759df7003e4 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -12,6 +12,7 @@ from pathlib import Path from queue import SimpleQueue import struct from typing import Final, cast +import uuid import zlib import httpx @@ -79,10 +80,15 @@ def _aws_str_header(name: str, value: str) -> bytes: ) -def _aws_event_frame(event_type: str, payload: Mapping[str, JsonValue], scenario_id: str) -> bytes: +def _aws_event_frame( + event_type: str, + payload: Mapping[str, JsonValue], + scenario_id: str, + unique_id: str, +) -> bytes: payload_bytes: Final = json.dumps(payload, separators=(",", ":")).replace( "$REQUEST_ID", scenario_id - ).encode() + ).replace("$UNIQUE_ID", unique_id).encode() headers_bytes: Final = ( _aws_str_header(":event-type", event_type) + _aws_str_header(":content-type", "application/json") @@ -209,11 +215,14 @@ class Provider: @staticmethod def _response(response: StoredResponse, scenario_id: str) -> Response: + unique_id: Final = f"{scenario_id}-{uuid.uuid4().hex[:8]}" match response: case JsonResponse(): return Response( content=json.dumps(response.body, separators=(",", ":")).replace( "$REQUEST_ID", scenario_id + ).replace( + "$UNIQUE_ID", unique_id ).encode(), media_type=response.content_type, status_code=response.status, @@ -227,13 +236,15 @@ class Provider: if response.frame_delay_ms > 0: async def stream() -> AsyncIterator[bytes]: for frame in response.frames: - yield f"{frame.replace('$REQUEST_ID', scenario_id)}\n\n".encode() + yield ( + f"{frame.replace('$REQUEST_ID', scenario_id).replace('$UNIQUE_ID', unique_id)}\n\n" + ).encode() await asyncio.sleep(response.frame_delay_ms / 1000) return StreamingResponse(stream(), media_type=response.content_type) stream_body: Final = ("\n\n".join(response.frames) + "\n\n").replace( "$REQUEST_ID", scenario_id - ) + ).replace("$UNIQUE_ID", unique_id) return Response(content=stream_body.encode(), media_type=response.content_type) case EventStreamResponse(): events: Final = ( @@ -244,6 +255,7 @@ class Provider: "bytes": base64.b64encode( json.dumps(event.payload, separators=(",", ":")) .replace("$REQUEST_ID", scenario_id) + .replace("$UNIQUE_ID", unique_id) .encode() ).decode(), }, @@ -254,7 +266,7 @@ class Provider: else response.events ) event_body: Final = b"".join( - _aws_event_frame(event.event_type, event.payload, scenario_id) for event in events + _aws_event_frame(event.event_type, event.payload, scenario_id, unique_id) for event in events ) return Response(content=event_body, media_type=response.content_type) diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index a70cd3619ae..d7f817efc3a 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -62,12 +62,12 @@ class FailureRow(BaseModel): class DailySpend(BaseModel): - model_config = ConfigDict(extra="ignore") + model_config = ConfigDict(frozen=True, extra="forbid") - spend: float | None = None - prompt_tokens: int | None = None - completion_tokens: int | None = None - api_requests: int | None = None + spend: float + prompt_tokens: int + completion_tokens: int + api_requests: int class Rollups(BaseModel): @@ -245,7 +245,6 @@ def register_scenario_deployment( *, response: StoredResponse | None = None, marker_suffix: str = "", - model_name: str | None = None, ) -> RegisteredDeployment: control_url: Final = os.environ["INTEGRATION_UPSTREAM_URL"].rstrip("/") run_marker: Final = sha256(key.encode()).hexdigest()[:12] @@ -254,7 +253,7 @@ def register_scenario_deployment( case.response if response is None else response, ) scenario.cleanups.callback(delete_scenario, handle) - registered_model_name: Final = model_name or f"cost-{marker}{marker_suffix}-{run_marker}" + registered_model_name: Final = f"cost-{marker}{marker_suffix}-{run_marker}" parameters: Final = { "model": case.litellm_model, "api_key": case.api_key, diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index f111778b233..effb6f2ed35 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -177,6 +177,7 @@ class RecountExpected(BaseModel): recount: RecountRates prompt_tokens: int | None = None completion_tokens: int | None = None + min_completion_tokens: int | None = None class FailureDetails(BaseModel): @@ -463,6 +464,23 @@ def data_errors() -> tuple[str, ...]: or not isinstance(case.expected, RecountExpected) ) ) + invalid_rollup_ids: Final = sorted( + case.name + for case in CASES + if isinstance(case.expected, ExactExpected) + and case.expected.rollups + and "$UNIQUE_ID" not in case.response.model_dump_json() + ) + invalid_pinned_tool_ids: Final = sorted( + case.name + for case in CASES + if isinstance(case.expected, RecountExpected) + and (case.expected.prompt_tokens is not None or case.expected.completion_tokens is not None) + and any( + marker in case.response.model_dump_json() + for marker in ('"id": "call_$REQUEST_ID"', '"id": "toolu_$REQUEST_ID"') + ) + ) return tuple( message for message in ( @@ -478,6 +496,10 @@ def data_errors() -> tuple[str, ...]: f"invalid passthrough opt-outs: {invalid_opt_outs}" if invalid_opt_outs else None, f"invalid fallback responses: {invalid_fallbacks}" if invalid_fallbacks else None, f"invalid disconnect cases: {invalid_disconnects}" if invalid_disconnects else None, + f"rollup responses lack $UNIQUE_ID: {invalid_rollup_ids}" if invalid_rollup_ids else None, + f"pinned tool IDs contain $REQUEST_ID: {invalid_pinned_tool_ids}" + if invalid_pinned_tool_ids + else None, ) if message is not None ) diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index 9fb7c6990a7..facab77828c 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -6930,7 +6930,7 @@ "response": { "content_type": "application/json", "body": { - "id": "msg_$REQUEST_ID", + "id": "msg_$UNIQUE_ID", "type": "message", "role": "assistant", "model": "claude-sonnet-5", @@ -7690,7 +7690,7 @@ "content_type": "text/event-stream", "frames": [ "event: message_start\ndata: {\"type\": \"message_start\", \"message\": {\"id\": \"msg_$REQUEST_ID\", \"type\": \"message\", \"role\": \"assistant\", \"model\": \"claude-sonnet-5\", \"content\": [], \"stop_reason\": null}}", - "event: content_block_start\ndata: {\"type\": \"content_block_start\", \"index\": 0, \"content_block\": {\"type\": \"tool_use\", \"id\": \"toolu_$REQUEST_ID\", \"name\": \"get_weather\", \"input\": {}}}", + "event: content_block_start\ndata: {\"type\": \"content_block_start\", \"index\": 0, \"content_block\": {\"type\": \"tool_use\", \"id\": \"call_fixture_0001\", \"name\": \"get_weather\", \"input\": {}}}", "event: content_block_delta\ndata: {\"type\": \"content_block_delta\", \"index\": 0, \"delta\": {\"type\": \"input_json_delta\", \"partial_json\": \"{\\\"city\\\": \\\"Berlin\\\", \\\"days\\\": 7, \\\"units\\\": \\\"metric\\\", \\\"notes\\\": \\\"filler filler filler filler fil\"}}", "event: content_block_delta\ndata: {\"type\": \"content_block_delta\", \"index\": 0, \"delta\": {\"type\": \"input_json_delta\", \"partial_json\": \"ler filler filler filler filler filler filler filler filler filler filler filler filler fi\"}}", "event: content_block_delta\ndata: {\"type\": \"content_block_delta\", \"index\": 0, \"delta\": {\"type\": \"input_json_delta\", \"partial_json\": \"ller filler filler filler filler filler filler filler filler filler filler filler filler \\\"}\"}}", @@ -7704,8 +7704,7 @@ "input_cost_per_token": 3e-06, "output_cost_per_token": 1.5e-05 }, - "prompt_tokens": 49, - "completion_tokens": 111 + "min_completion_tokens": 60 } }, { @@ -12877,8 +12876,7 @@ "input_cost_per_token": 5.2e-07, "output_cost_per_token": 3.12e-06 }, - "prompt_tokens": 46, - "completion_tokens": 88 + "min_completion_tokens": 60 } }, { @@ -20752,7 +20750,7 @@ "response": { "content_type": "application/json", "body": { - "id": "chatcmpl-$REQUEST_ID", + "id": "chatcmpl-$UNIQUE_ID", "object": "chat.completion", "created": 1789788262, "model": "gpt-5.6", @@ -21610,7 +21608,7 @@ "content_type": "text/event-stream", "frames": [ "data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"role\": \"assistant\"}, \"finish_reason\": null}], \"usage\": null}", - "data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"role\": \"assistant\", \"tool_calls\": [{\"index\": 0, \"id\": \"call_$REQUEST_ID\", \"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"arguments\": \"\"}}]}, \"finish_reason\": null}], \"usage\": null}", + "data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"role\": \"assistant\", \"tool_calls\": [{\"index\": 0, \"id\": \"call_fixture_0001\", \"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"arguments\": \"\"}}]}, \"finish_reason\": null}], \"usage\": null}", "data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"tool_calls\": [{\"index\": 0, \"function\": {\"arguments\": \"{\\\"city\\\": \\\"Berlin\\\", \\\"days\\\": 7, \\\"units\\\": \\\"metric\\\", \\\"notes\\\": \\\"filler filler filler filler fil\"}}]}, \"finish_reason\": null}], \"usage\": null}", "data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"tool_calls\": [{\"index\": 0, \"function\": {\"arguments\": \"ler filler filler filler filler filler filler filler filler filler filler filler filler fi\"}}]}, \"finish_reason\": null}], \"usage\": null}", "data: {\"id\": \"chatcmpl-$REQUEST_ID\", \"object\": \"chat.completion.chunk\", \"created\": 1789788263, \"model\": \"gpt-5.6\", \"choices\": [{\"index\": 0, \"delta\": {\"tool_calls\": [{\"index\": 0, \"function\": {\"arguments\": \"ller filler filler filler filler filler filler filler filler filler filler filler filler \\\"}\"}}]}, \"finish_reason\": null}], \"usage\": null}", @@ -21623,8 +21621,7 @@ "input_cost_per_token": 1.75e-06, "output_cost_per_token": 1.4e-05 }, - "prompt_tokens": 44, - "completion_tokens": 105 + "min_completion_tokens": 60 } }, { diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index e8eed98205e..15b5921c94e 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -209,6 +209,35 @@ def _assert_exact( assert_total_is_sum_of_components(row, breakdown, case.name) +def _assert_recount(case: CostTrackingTestCase, expected: RecountExpected, row: CostRow) -> None: + assert row.prompt_tokens is not None and row.prompt_tokens > 0, ( + f"{case.name}: recount case counted no input tokens: prompt_tokens={row.prompt_tokens}" + ) + assert row.completion_tokens is not None and row.completion_tokens > 0, ( + f"{case.name}: recount case counted no output tokens: completion_tokens={row.completion_tokens}" + ) + if expected.prompt_tokens is not None: + assert row.prompt_tokens == expected.prompt_tokens, ( + f"{case.name}: prompt_tokens {row.prompt_tokens} != pinned {expected.prompt_tokens}" + ) + if expected.completion_tokens is not None: + assert row.completion_tokens == expected.completion_tokens, ( + f"{case.name}: completion_tokens {row.completion_tokens} != pinned {expected.completion_tokens}" + ) + if expected.min_completion_tokens is not None: + assert row.completion_tokens >= expected.min_completion_tokens, ( + f"{case.name}: completion_tokens {row.completion_tokens} < minimum {expected.min_completion_tokens}" + ) + recount: Final = row.prompt_tokens * expected.recount.input_cost_per_token + ( + row.completion_tokens * expected.recount.output_cost_per_token + ) + assert row.spend is not None and approx_equal(row.spend, recount), ( + f"{case.name}: spend {row.spend} != recount {recount} at map rates" + ) + assert row.breakdown is not None, f"{case.name}: no cost_breakdown persisted" + assert_total_is_sum_of_components(row, row.breakdown, case.name) + + @pytest.mark.parametrize("case", _CASES) def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) -> None: marker: Final = sha256(case.name.encode()).hexdigest()[:12] @@ -251,20 +280,6 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) if case.fallback_from is not None else None ) - if isinstance(expected, ExactExpected) and expected.rollups: - assert deployment is not None - rollup_deployments: Final = tuple( - register_scenario_deployment( - scenario, - case, - marker, - key, - marker_suffix=f"-r{index}", - model_name=deployment.model_name, - ) - for index in (2, 3) - ) - assert len(rollup_deployments) == 2 model_name: Final = ( case.model if passthrough_provider in {"gemini", "anthropic"} @@ -334,18 +349,8 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert len(frames) == case.disconnect_after_frames row: Final = poll_cost_row(key) assert isinstance(expected, RecountExpected) - if expected.prompt_tokens is not None: - assert row.prompt_tokens == expected.prompt_tokens - if expected.completion_tokens is not None: - assert row.completion_tokens == expected.completion_tokens - assert row.prompt_tokens is not None and row.prompt_tokens > 0 - assert row.completion_tokens is not None and row.completion_tokens > 0 - recount: Final = row.prompt_tokens * expected.recount.input_cost_per_token + ( - row.completion_tokens * expected.recount.output_cost_per_token - ) - assert row.spend is not None and approx_equal(row.spend, recount) - assert row.breakdown is not None - assert_total_is_sum_of_components(row, row.breakdown, case.name) + assert row.status == "success", f"{case.name}: disconnect row status was {row.status}" + _assert_recount(case, expected, row) return responses: Final = tuple( ( @@ -374,21 +379,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) rows: Final = poll_rows(key) if len(responses) > 1 else (poll_cost_row(key),) if isinstance(expected, RecountExpected): row: Final = rows[0] - assert row.prompt_tokens is not None and row.prompt_tokens > 0, ( - f"{case.name}: recount case counted no input tokens: prompt_tokens={row.prompt_tokens}" - ) - assert row.completion_tokens is not None and row.completion_tokens > 0, ( - f"{case.name}: recount case counted no output tokens: completion_tokens={row.completion_tokens}" - ) - recount: Final = row.prompt_tokens * case.expected.recount.input_cost_per_token + ( - row.completion_tokens * case.expected.recount.output_cost_per_token - ) - assert row.spend is not None and approx_equal(row.spend, recount), ( - f"{case.name}: spend {row.spend} != recount {recount} at map rates" - ) - 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) + _assert_recount(case, expected, row) return assert isinstance(expected, ExactExpected) if fallback_deployment is not None: @@ -423,8 +414,8 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert approx_equal(rollups.team_spend, target_spend) assert approx_equal(rollups.user_spend, target_spend) assert approx_equal(rollups.end_user_spend, target_spend) - assert approx_equal(rollups.daily_user.spend or 0.0, target_spend) - assert approx_equal(rollups.daily_team.spend or 0.0, target_spend) + assert approx_equal(rollups.daily_user.spend, target_spend) + assert approx_equal(rollups.daily_team.spend, target_spend) assert rollups.daily_user.prompt_tokens == expected.prompt_tokens * 3 assert rollups.daily_user.completion_tokens == expected.completion_tokens * 3 assert rollups.daily_user.api_requests == 3 From cc93a37322d1c8df2451654de5879b982eb5b60b Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 00:48:42 +0000 Subject: [PATCH 3/7] test(integration): wait for every rollup write and bound the disconnect recount Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integration/cost_calculation/conftest.py | 23 +++++++++++++++---- .../cost_calculation/cost_tracking_case.py | 1 + .../cost_calculation/cost_tracking_cases.json | 5 +++- .../cost_calculation/test_cost_tracking.py | 8 +++++-- 4 files changed, 29 insertions(+), 8 deletions(-) diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index d7f817efc3a..fb20f5a9cc2 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -123,7 +123,7 @@ def poll_cost_row(key: str) -> CostRow: return result -def poll_rows(key: str) -> tuple[CostRow, ...]: +def poll_rows(key: str, count: int) -> tuple[CostRow, ...]: digest: Final = sha256(key.encode()).hexdigest() def read() -> tuple[CostRow, ...]: @@ -134,11 +134,11 @@ def poll_rows(key: str) -> tuple[CostRow, ...]: ) return tuple(parsed for row in rows if (parsed := _row(row)) is not None) - result: Final = eventually(read, lambda rows: len(rows) > 0, seconds=20) + result: Final = eventually(read, lambda rows: len(rows) >= count, seconds=60) return result -def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str) -> Rollups: +def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str, requests: int, spend: float) -> Rollups: digest: Final = sha256(key.encode()).hexdigest() def read() -> Rollups | None: @@ -170,7 +170,7 @@ def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str) -> Roll ) if not all((key_rows, team_rows, user_rows, end_user_rows, daily_user_rows, daily_team_rows)): return None - return Rollups( + rollups: Final = Rollups( key_spend=float(key_rows[0]["spend"]), team_spend=float(team_rows[0]["spend"]), user_spend=float(user_rows[0]["spend"]), @@ -178,8 +178,21 @@ def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str) -> Roll daily_user=DailySpend.model_validate(daily_user_rows[0]), daily_team=DailySpend.model_validate(daily_team_rows[0]), ) + if rollups.daily_user.api_requests < requests or rollups.daily_team.api_requests < requests: + return None + if not all( + approx_equal(actual, spend) + for actual in ( + rollups.key_spend, + rollups.team_spend, + rollups.user_spend, + rollups.end_user_spend, + ) + ): + return None + return rollups - result: Final = eventually(read, lambda value: value is not None, seconds=20) + result: Final = eventually(read, lambda value: value is not None, seconds=60) assert result is not None return result diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index effb6f2ed35..56a7cefb443 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -178,6 +178,7 @@ class RecountExpected(BaseModel): prompt_tokens: int | None = None completion_tokens: int | None = None min_completion_tokens: int | None = None + max_completion_tokens: int | None = None class FailureDetails(BaseModel): diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index facab77828c..78d00f28381 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -30014,7 +30014,10 @@ "recount": { "input_cost_per_token": 1.75e-06, "output_cost_per_token": 1.4e-05 - } + }, + "prompt_tokens": 10, + "min_completion_tokens": 9, + "max_completion_tokens": 30 } } ] diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index 15b5921c94e..692b76fff15 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -228,6 +228,10 @@ def _assert_recount(case: CostTrackingTestCase, expected: RecountExpected, row: assert row.completion_tokens >= expected.min_completion_tokens, ( f"{case.name}: completion_tokens {row.completion_tokens} < minimum {expected.min_completion_tokens}" ) + if expected.max_completion_tokens is not None: + assert row.completion_tokens <= expected.max_completion_tokens, ( + f"{case.name}: completion_tokens {row.completion_tokens} > maximum {expected.max_completion_tokens}" + ) recount: Final = row.prompt_tokens * expected.recount.input_cost_per_token + ( row.completion_tokens * expected.recount.output_cost_per_token ) @@ -376,7 +380,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert response.is_success, f"{case.name}: proxy returned {response.status_code}: {response.text[:400]}" if case.response.content_type == "text/event-stream": _assert_stream_has_no_error(response.text) - rows: Final = poll_rows(key) if len(responses) > 1 else (poll_cost_row(key),) + rows: Final = poll_rows(key, len(responses)) if isinstance(expected, RecountExpected): row: Final = rows[0] _assert_recount(case, expected, row) @@ -408,8 +412,8 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) if expected.rollups: assert deployment is not None and team_id is not None and user_id is not None assert end_user_id is not None - rollups: Final = poll_rollups(key, team_id, user_id, end_user_id) target_spend: Final = expected.spend * 3 + rollups: Final = poll_rollups(key, team_id, user_id, end_user_id, requests=3, spend=target_spend) assert approx_equal(rollups.key_spend, target_spend) assert approx_equal(rollups.team_spend, target_spend) assert approx_equal(rollups.user_spend, target_spend) From 8e3bb5daabf9d68e7646ddaf09c0678d96fcb0a2 Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 00:49:51 +0000 Subject: [PATCH 4/7] test(integration): wait for all rollup writes and pin fallback and disconnect rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integration/cost_calculation/conftest.py | 24 +++++++++++++++---- .../cost_calculation/cost_tracking_case.py | 1 + .../cost_calculation/cost_tracking_cases.json | 4 +++- .../cost_calculation/test_cost_tracking.py | 12 ++++++++-- 4 files changed, 34 insertions(+), 7 deletions(-) diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index d7f817efc3a..4f1f723c2f7 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -123,7 +123,7 @@ def poll_cost_row(key: str) -> CostRow: return result -def poll_rows(key: str) -> tuple[CostRow, ...]: +def poll_rows(key: str, count: int) -> tuple[CostRow, ...]: digest: Final = sha256(key.encode()).hexdigest() def read() -> tuple[CostRow, ...]: @@ -134,11 +134,18 @@ def poll_rows(key: str) -> tuple[CostRow, ...]: ) return tuple(parsed for row in rows if (parsed := _row(row)) is not None) - result: Final = eventually(read, lambda rows: len(rows) > 0, seconds=20) + result: Final = eventually(read, lambda rows: len(rows) >= count, seconds=60) return result -def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str) -> Rollups: +def poll_rollups( + key: str, + team_id: str, + user_id: str, + end_user_id: str, + requests: int, + target_spend: float, +) -> Rollups: digest: Final = sha256(key.encode()).hexdigest() def read() -> Rollups | None: @@ -179,7 +186,16 @@ def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str) -> Roll daily_team=DailySpend.model_validate(daily_team_rows[0]), ) - result: Final = eventually(read, lambda value: value is not None, seconds=20) + result: Final = eventually( + read, + lambda value: ( + value is not None + and value.daily_user.api_requests >= requests + and value.daily_team.api_requests >= requests + and approx_equal(value.key_spend, target_spend) + ), + seconds=60, + ) assert result is not None return result diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index effb6f2ed35..56a7cefb443 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -178,6 +178,7 @@ class RecountExpected(BaseModel): prompt_tokens: int | None = None completion_tokens: int | None = None min_completion_tokens: int | None = None + max_completion_tokens: int | None = None class FailureDetails(BaseModel): diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index facab77828c..bc5b40ccafe 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -30014,7 +30014,9 @@ "recount": { "input_cost_per_token": 1.75e-06, "output_cost_per_token": 1.4e-05 - } + }, + "prompt_tokens": 8, + "max_completion_tokens": 30 } } ] diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index 15b5921c94e..de55ab396fe 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -228,6 +228,10 @@ def _assert_recount(case: CostTrackingTestCase, expected: RecountExpected, row: assert row.completion_tokens >= expected.min_completion_tokens, ( f"{case.name}: completion_tokens {row.completion_tokens} < minimum {expected.min_completion_tokens}" ) + if expected.max_completion_tokens is not None: + assert row.completion_tokens <= expected.max_completion_tokens, ( + f"{case.name}: completion_tokens {row.completion_tokens} > maximum {expected.max_completion_tokens}" + ) recount: Final = row.prompt_tokens * expected.recount.input_cost_per_token + ( row.completion_tokens * expected.recount.output_cost_per_token ) @@ -376,7 +380,11 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert response.is_success, f"{case.name}: proxy returned {response.status_code}: {response.text[:400]}" if case.response.content_type == "text/event-stream": _assert_stream_has_no_error(response.text) - rows: Final = poll_rows(key) if len(responses) > 1 else (poll_cost_row(key),) + rows: Final = ( + poll_rows(key, len(responses)) + if len(responses) > 1 or fallback_deployment is not None + else (poll_cost_row(key),) + ) if isinstance(expected, RecountExpected): row: Final = rows[0] _assert_recount(case, expected, row) @@ -408,8 +416,8 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) if expected.rollups: assert deployment is not None and team_id is not None and user_id is not None assert end_user_id is not None - rollups: Final = poll_rollups(key, team_id, user_id, end_user_id) target_spend: Final = expected.spend * 3 + rollups: Final = poll_rollups(key, team_id, user_id, end_user_id, 3, target_spend) assert approx_equal(rollups.key_spend, target_spend) assert approx_equal(rollups.team_spend, target_spend) assert approx_equal(rollups.user_spend, target_spend) From dc9889a4813a8979ddf61334ac30f570317cc124 Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 00:51:08 +0000 Subject: [PATCH 5/7] test(integration): restore the concurrent rollup and disconnect fix from cc93a37 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integration/cost_calculation/conftest.py | 32 ++++++++----------- .../cost_calculation/cost_tracking_cases.json | 3 +- .../cost_calculation/test_cost_tracking.py | 8 ++--- 3 files changed, 18 insertions(+), 25 deletions(-) diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index 3536c020515..fb20f5a9cc2 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -138,14 +138,7 @@ def poll_rows(key: str, count: int) -> tuple[CostRow, ...]: return result -def poll_rollups( - key: str, - team_id: str, - user_id: str, - end_user_id: str, - requests: int, - target_spend: float, -) -> Rollups: +def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str, requests: int, spend: float) -> Rollups: digest: Final = sha256(key.encode()).hexdigest() def read() -> Rollups | None: @@ -185,18 +178,21 @@ def poll_rollups( daily_user=DailySpend.model_validate(daily_user_rows[0]), daily_team=DailySpend.model_validate(daily_team_rows[0]), ) + if rollups.daily_user.api_requests < requests or rollups.daily_team.api_requests < requests: + return None + if not all( + approx_equal(actual, spend) + for actual in ( + rollups.key_spend, + rollups.team_spend, + rollups.user_spend, + rollups.end_user_spend, + ) + ): + return None return rollups - result: Final = eventually( - read, - lambda value: ( - value is not None - and value.daily_user.api_requests >= requests - and value.daily_team.api_requests >= requests - and approx_equal(value.key_spend, target_spend) - ), - seconds=60, - ) + result: Final = eventually(read, lambda value: value is not None, seconds=60) assert result is not None return result diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index bc5b40ccafe..78d00f28381 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -30015,7 +30015,8 @@ "input_cost_per_token": 1.75e-06, "output_cost_per_token": 1.4e-05 }, - "prompt_tokens": 8, + "prompt_tokens": 10, + "min_completion_tokens": 9, "max_completion_tokens": 30 } } diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index de55ab396fe..692b76fff15 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -380,11 +380,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert response.is_success, f"{case.name}: proxy returned {response.status_code}: {response.text[:400]}" if case.response.content_type == "text/event-stream": _assert_stream_has_no_error(response.text) - rows: Final = ( - poll_rows(key, len(responses)) - if len(responses) > 1 or fallback_deployment is not None - else (poll_cost_row(key),) - ) + rows: Final = poll_rows(key, len(responses)) if isinstance(expected, RecountExpected): row: Final = rows[0] _assert_recount(case, expected, row) @@ -417,7 +413,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert deployment is not None and team_id is not None and user_id is not None assert end_user_id is not None target_spend: Final = expected.spend * 3 - rollups: Final = poll_rollups(key, team_id, user_id, end_user_id, 3, target_spend) + rollups: Final = poll_rollups(key, team_id, user_id, end_user_id, requests=3, spend=target_spend) assert approx_equal(rollups.key_spend, target_spend) assert approx_equal(rollups.team_spend, target_spend) assert approx_equal(rollups.user_spend, target_spend) From 6e77d23f4d50503da2b4d35ab883df89744e14f8 Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 00:52:25 +0000 Subject: [PATCH 6/7] test(integration): settle rollup and fallback row polling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/_support/client.py | 9 ++- .../integration/cost_calculation/conftest.py | 64 ++++++++++++------- .../cost_calculation/test_cost_tracking.py | 20 ++++-- 3 files changed, 65 insertions(+), 28 deletions(-) diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 0b6771623c0..e07cbe6b2a3 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -34,12 +34,19 @@ def delete_key_if_present(candidate: Gateway, key: str) -> None: assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [] -def eventually(read: Callable[[], T], satisfied: Callable[[T], bool], seconds: float = 10) -> T: +def eventually( + read: Callable[[], T], + satisfied: Callable[[T], bool], + seconds: float = 10, + return_last_on_timeout: bool = False, +) -> T: deadline: Final = time.monotonic() + seconds while True: observed: Final = read() if satisfied(observed): return observed + if return_last_on_timeout and time.monotonic() >= deadline: + return observed assert time.monotonic() < deadline, f"State did not converge: {observed!r}" time.sleep(0.1) diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index fb20f5a9cc2..bc68419f01d 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -123,22 +123,33 @@ def poll_cost_row(key: str) -> CostRow: return result -def poll_rows(key: str, count: int) -> tuple[CostRow, ...]: +def read_rows_now(key: str) -> tuple[CostRow, ...]: digest: Final = sha256(key.encode()).hexdigest() + rows: Final = read_rows( + 'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id ' + 'FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"', + (digest,), + ) + return tuple(parsed for row in rows if (parsed := _row(row)) is not None) - def read() -> tuple[CostRow, ...]: - rows: Final = read_rows( - 'SELECT spend, status, metadata, prompt_tokens, completion_tokens, model_id ' - 'FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"', - (digest,), - ) - return tuple(parsed for row in rows if (parsed := _row(row)) is not None) - result: Final = eventually(read, lambda rows: len(rows) >= count, seconds=60) +def poll_rows(key: str, count: int) -> tuple[CostRow, ...]: + result: Final = eventually( + lambda: read_rows_now(key), + lambda rows: len(rows) >= count, + seconds=60, + ) return result -def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str, requests: int, spend: float) -> Rollups: +def poll_rollups( + key: str, + team_id: str, + user_id: str, + end_user_id: str, + target_spend: float, + target_requests: int, +) -> Rollups: digest: Final = sha256(key.encode()).hexdigest() def read() -> Rollups | None: @@ -178,21 +189,28 @@ def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str, request daily_user=DailySpend.model_validate(daily_user_rows[0]), daily_team=DailySpend.model_validate(daily_team_rows[0]), ) - if rollups.daily_user.api_requests < requests or rollups.daily_team.api_requests < requests: - return None - if not all( - approx_equal(actual, spend) - for actual in ( - rollups.key_spend, - rollups.team_spend, - rollups.user_spend, - rollups.end_user_spend, - ) - ): - return None return rollups - result: Final = eventually(read, lambda value: value is not None, seconds=60) + def settled(value: Rollups | None) -> bool: + return value is not None and all( + ( + approx_equal(value.key_spend, target_spend), + approx_equal(value.team_spend, target_spend), + approx_equal(value.user_spend, target_spend), + approx_equal(value.end_user_spend, target_spend), + approx_equal(value.daily_user.spend, target_spend), + approx_equal(value.daily_team.spend, target_spend), + value.daily_user.api_requests == target_requests, + value.daily_team.api_requests == target_requests, + ) + ) + + result: Final = eventually( + read, + settled, + seconds=20, + return_last_on_timeout=True, + ) assert result is not None return result diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index 692b76fff15..05346396a8a 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -7,6 +7,7 @@ from itertools import islice import json from hashlib import sha256 import struct +import time from typing import Final, cast import uuid import wave @@ -27,6 +28,7 @@ from integration.cost_calculation.conftest import ( poll_failure_row, poll_rollups, poll_rows, + read_rows_now, register_scenario_deployment, ) from integration.cost_calculation.cost_tracking_case import ( @@ -388,9 +390,11 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert isinstance(expected, ExactExpected) if fallback_deployment is not None: assert deployment is not None - assert len(rows) == 1 - assert rows[0].status == "success" - assert rows[0].model_id == deployment.identity + time.sleep(3) + settled_rows: Final = read_rows_now(key) + assert len(settled_rows) == 1 + assert settled_rows[0].status == "success" + assert settled_rows[0].model_id == deployment.identity if isinstance(case.response, BinaryResponse): header: Final = response.headers.get("x-litellm-response-cost") if header is not None: @@ -413,7 +417,15 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert deployment is not None and team_id is not None and user_id is not None assert end_user_id is not None target_spend: Final = expected.spend * 3 - rollups: Final = poll_rollups(key, team_id, user_id, end_user_id, requests=3, spend=target_spend) + target_requests: Final = 3 + rollups: Final = poll_rollups( + key, + team_id, + user_id, + end_user_id, + target_spend, + target_requests, + ) assert approx_equal(rollups.key_spend, target_spend) assert approx_equal(rollups.team_spend, target_spend) assert approx_equal(rollups.user_spend, target_spend) From 96ca550377393e8e0077e214560d1a0aef991385 Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 01:00:25 +0000 Subject: [PATCH 7/7] test(integration): register deployments for bedrock passthrough cases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/cost_calculation/test_cost_tracking.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index 05346396a8a..f2057067c67 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -271,7 +271,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) scenario.cleanups.callback(delete_scenario, scenario_handle) deployment: Final = ( register_scenario_deployment(scenario, case, marker, key) - if passthrough_provider is None + if passthrough_provider not in {"gemini", "anthropic"} else None ) fallback_deployment: Final = (