From 2edea0be086ebbf9c7c6ae0c1a539b64c588fabe Mon Sep 17 00:00:00 2001 From: kerry Date: Sun, 20 Sep 2026 00:36:18 +0000 Subject: [PATCH] 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