From 80d804d6f98b69ff280b2020df76cc3bd6aa07fe Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 14 Sep 2026 22:26:20 -0700 Subject: [PATCH] test(spend): preserve multi-day coverage and immutable assertions --- .../spend_tracking/spend_reconciliation.py | 12 +++-- .../spend_tracking/test_spend_tracking_e2e.py | 13 +++-- .../test_team_daily_activity_e2e.py | 48 +++++++++++++------ 3 files changed, 51 insertions(+), 22 deletions(-) diff --git a/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py b/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py index 8fcfea3f296..26809874aed 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py +++ b/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py @@ -95,9 +95,10 @@ def assert_logs_match(client: SpendClient, traffic: TeamTraffic) -> None: assert frozenset(row.request_id for row in rows) == expected_ids, "stored IDs must equal returned response IDs" assert len(rows) == len(traffic.responses), "expected exactly one scoped spend row per response" by_id: Final = {row.request_id: row for row in rows} - for response in traffic.responses: - row = by_id[response.id] - usage = response.usage + + def assert_response(response: ChatResponse) -> None: + row: Final = by_id[response.id] + usage: Final = response.usage assert usage is not None and usage.prompt_tokens is not None and usage.completion_tokens is not None assert row.team_id == traffic.team_id assert row.status == "success" @@ -105,5 +106,8 @@ def assert_logs_match(client: SpendClient, traffic: TeamTraffic) -> None: assert row.prompt_tokens == usage.prompt_tokens assert row.completion_tokens == usage.completion_tokens assert row.total_tokens == usage.total_tokens - expected_cost = usage.prompt_tokens * INPUT_RATE + usage.completion_tokens * OUTPUT_RATE + expected_cost: Final = usage.prompt_tokens * INPUT_RATE + usage.completion_tokens * OUTPUT_RATE assert row.spend is not None and isclose(row.spend, expected_cost, rel_tol=1e-6, abs_tol=1e-9) + + for response in traffic.responses: + assert_response(response) diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py index 750285cbfb6..8a91e53e7d7 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py @@ -18,6 +18,7 @@ fails the test; a pricing or token-count drift does not. import time from collections.abc import Callable from math import isclose +from typing import Final import pytest from e2e_http import Success @@ -284,14 +285,18 @@ def test_key_spend_equals_sum_of_logs(client: SpendClient, scoped_key: str) -> N def test_burst_of_concurrent_calls_loses_no_spend( client: SpendClient, resources: ResourceManager ) -> None: - from spend_reconciliation import assert_logs_match, create_traffic + from spend_reconciliation import TeamTraffic, assert_logs_match, create_traffic - traffic = create_traffic(client, resources) - for team in traffic: + traffic: Final = create_traffic(client, resources) + + def assert_team(team: TeamTraffic) -> None: assert_logs_match(client, team) - key_spend = client.poll_key_spend(team.key, minimum=team.spend * 0.999999) + key_spend: Final = client.poll_key_spend(team.key, minimum=team.spend * 0.999999) assert isclose(key_spend, team.spend, rel_tol=1e-6, abs_tol=1e-9) + for team in traffic: + assert_team(team) + @pytest.mark.covers("quota_management.spend_tracking.pagination.keeps_total") def test_spend_logs_v2_pagination_caps_pages_and_keeps_total( diff --git a/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py b/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py index dca5572f510..ef635e59743 100644 --- a/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py @@ -15,9 +15,10 @@ from typing import Final import pytest from e2e_http import ProbeResult from lifecycle import ResourceManager +from proxy_client import Converged, await_converged from pydantic import BaseModel from spend_e2e_client import SpendClient -from spend_reconciliation import assert_logs_match, create_traffic +from spend_reconciliation import TeamTraffic, assert_logs_match, create_traffic pytestmark = pytest.mark.e2e @@ -92,7 +93,7 @@ class TestTeamDailyActivity: team_ids: Final = ",".join(team.team_id for team in traffic) def fetch( - page: int, start: str = started.isoformat(), end: str = ended.isoformat() + page: int, start: str = (started - timedelta(days=1)).isoformat(), end: str = ended.isoformat() ) -> TeamDailyActivityResponse: result: Final = _probe( client, @@ -112,22 +113,27 @@ class TestTeamDailyActivity: assert first.metadata.total_pages <= len(traffic) * 2, "unexpected extra scoped daily groups" return (first, *(fetch(page) for page in range(2, first.metadata.total_pages + 1))) - deadline: Final = time.monotonic() + client.proxy.poll_timeout - while True: - observed = pages() - if sum(page.metadata.total_api_requests for page in observed) >= sum(len(t.responses) for t in traffic): - break - if time.monotonic() >= deadline: - break - time.sleep(client.proxy.poll_interval) + outcome: Final = await_converged( + pages, + converged=lambda values: ( + sum(page.metadata.total_api_requests for page in values) >= sum(len(team.responses) for team in traffic) + ), + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + observed: Final = outcome.result if isinstance(outcome, Converged) else outcome.last_result + assert observed is not None, "daily aggregation must return a response before the deadline" assert len(observed) >= 2, "two teams must exercise a page boundary" - for index, page in enumerate(observed, 1): + + def assert_page(index: int, page: TeamDailyActivityResponse) -> None: assert page.metadata.page == index assert page.metadata.total_pages == len(observed) assert page.metadata.has_more == (index < len(observed)) assert len(page.results) == 1, "each fetched daily group must appear in results" - row = page.results[0] + row: Final = page.results[0] assert started <= datetime.fromisoformat(row.date).date() <= ended assert len(row.breakdown.entities) == 1 assert row.metrics.total_tokens == page.metadata.total_tokens @@ -138,6 +144,9 @@ class TestTeamDailyActivity: assert row.metrics.failed_requests == page.metadata.total_failed_requests assert isclose(row.metrics.spend, page.metadata.total_spend, rel_tol=1e-6, abs_tol=1e-9) + for index, page in enumerate(observed, 1): + assert_page(index, page) + entities: Final = tuple( (team_id, entity.metrics) for page in observed @@ -145,8 +154,9 @@ class TestTeamDailyActivity: for team_id, entity in row.breakdown.entities.items() ) assert frozenset(team_id for team_id, _ in entities) == frozenset(team.team_id for team in traffic) - for team in traffic: - metrics = tuple(metrics for team_id, metrics in entities if team_id == team.team_id) + + def assert_team(team: TeamTraffic) -> None: + metrics: Final = tuple(metrics for team_id, metrics in entities if team_id == team.team_id) assert sum(m.api_requests for m in metrics) == len(team.responses) assert sum(m.successful_requests for m in metrics) == len(team.responses) assert sum(m.failed_requests for m in metrics) == 0 @@ -154,6 +164,10 @@ class TestTeamDailyActivity: assert sum(m.completion_tokens for m in metrics) == team.completion_tokens assert sum(m.total_tokens for m in metrics) == team.prompt_tokens + team.completion_tokens assert isclose(sum(m.spend for m in metrics), team.spend, rel_tol=1e-6, abs_tol=1e-9) + + for team in traffic: + assert_team(team) + assert isclose( sum(page.metadata.total_spend for page in observed), sum(team.spend for team in traffic), @@ -164,6 +178,12 @@ class TestTeamDailyActivity: team.prompt_tokens + team.completion_tokens for team in traffic ) + for days in (7, 30): + assert ( + tuple(fetch(page, (started - timedelta(days=days)).isoformat()) for page in range(1, len(observed) + 1)) + == observed + ), f"{days}-day activity must preserve the same isolated groups and totals" + empty_date: Final = (started - timedelta(days=7)).isoformat() empty: Final = fetch(1, empty_date, empty_date) assert empty.results == []