From 0e84aa5b5ac3e6699bf1e43f150702de912a59c3 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 23 Jun 2026 18:52:34 -0700 Subject: [PATCH] test(e2e): drop flaky test_tag_spend_matches_sum_of_tagged_logs The test wrote tagged requests and polled /spend/tags expecting read-after-write consistency. /spend/tags itself is fine; verified live that request_tags is stored as a JSON array and the endpoint reflects a fresh tag within seconds, so the failures were a timing flake under full-suite load rather than a real defect. Coverage is retained by test_request_tags_round_trip (tags persist onto the row) and the /spend/tags route probe in test_spend_routes.py. Also remove the now-dead tag-spend scaffolding this test was the only user of: poll_tag_spend, spend_by_tags, TagSpendPoll, SpendTagsError, the TagSpend/TagSpends models, and their imports. --- tests/e2e/models.py | 23 ++--- .../SPEND_TRACKING_COVERAGE_MATRIX.md | 7 +- tests/e2e/spend_tracking/spend_e2e_client.py | 84 +---------------- .../spend_tracking/test_spend_tracking_e2e.py | 93 +++++++------------ 4 files changed, 44 insertions(+), 163 deletions(-) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index e90969da6ea..fbeb3d44fa5 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -166,19 +166,6 @@ class SpendCalculateResponse(BaseModel): cost: float -# ---------- spend tags ---------- - - -class TagSpend(BaseModel): - individual_request_tag: str - log_count: int | None = None - total_spend: float | None = None - - -class TagSpends(RootModel[list[TagSpend]]): - pass - - # ---------- route probing ---------- @@ -228,10 +215,14 @@ class CustomPricing(BaseModel): def token_cost(self, prompt_tokens: int, completion_tokens: int) -> float: """Spend for a fresh (uncached) call under these rates: the proxy's custom-pricing formula (prompt * input + completion * output).""" - assert self.input_cost_per_token is not None and self.output_cost_per_token is not None, ( - "custom pricing has no per-token rates" + assert ( + self.input_cost_per_token is not None + and self.output_cost_per_token is not None + ), "custom pricing has no per-token rates" + return ( + prompt_tokens * self.input_cost_per_token + + completion_tokens * self.output_cost_per_token ) - return prompt_tokens * self.input_cost_per_token + completion_tokens * self.output_cost_per_token class ModelInfoEntry(BaseModel): diff --git a/tests/e2e/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md b/tests/e2e/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md index fbb5ce60c51..0c6322c8251 100644 --- a/tests/e2e/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md +++ b/tests/e2e/spend_tracking/SPEND_TRACKING_COVERAGE_MATRIX.md @@ -39,9 +39,9 @@ proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`. | Entity | Existing | Status | Live e2e | |--------|----------|--------|----------| | API key | `test_db_spend_update_writer.py`, `test_spend_counters.py` | covered | yes (`test_key_spend_equals_sum_of_logs`) | -| Tag | `test_update_daily_tag_spend.py` | partial | yes (`test_tag_spend_matches_sum_of_tagged_logs`) | +| Tag | `test_update_daily_tag_spend.py` | partial | yes (`test_request_tags_round_trip`, propagation only) | | End-user | `test_proxy_update_spend.py` | covered | yes | -| Spend == sum(logs) consistency | none | gap | yes (key + tag aggregate == sum of rows) | +| Spend == sum(logs) consistency | none | gap | yes (key aggregate == sum of rows) | ## Spend read endpoints (verification surface) @@ -49,7 +49,7 @@ proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`. |----------|----------|--------|----------| | `/spend/logs` (request_id / api_key) | `test_spend_management_endpoints.py` | covered | yes (primary read path) | | `/spend/calculate` | `local_testing/test_spend_calculate_endpoint.py` | covered | yes (`test_spend_calculate_returns_nonzero_cost`) | -| `/spend/tags` | `test_spend_management_endpoints.py` | partial | yes (tag accuracy test) | +| `/spend/tags` | `test_spend_management_endpoints.py` | partial | yes (`test_spend_routes.py` route probe) | | whole spend GET surface (22 routes) | unit per-handler | partial | yes (`test_spend_routes.py` probes each for 404/5xx) | ## What this suite pins @@ -62,7 +62,6 @@ proxy + SpendLogs rows. Status: `covered` / `partial` / `gap`. | `test_cache_hit_is_zero_cost_and_suffixed` | cache hits not double-charged; `_cache_hit` suffix | | `test_key_spend_equals_sum_of_logs` | key aggregate == sum of rows | | `test_request_tags_round_trip` | tags persist onto the row | -| `test_tag_spend_matches_sum_of_tagged_logs` | `/spend/tags` SUM/COUNT == tagged rows | | `test_end_user_spend_attributed_on_row` | `end_user` attributed + costed | | `test_failure_call_writes_failure_status_row` | failed call -> `status=failure`, `spend=0` | | `test_spend_calculate_returns_nonzero_cost` | cost-map smoke (no batch wait) | diff --git a/tests/e2e/spend_tracking/spend_e2e_client.py b/tests/e2e/spend_tracking/spend_e2e_client.py index 1645297cad4..d749d69f1a4 100644 --- a/tests/e2e/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/spend_tracking/spend_e2e_client.py @@ -2,8 +2,8 @@ Generic proxy operations (keys, customers, chat/embed, route probing, SpendLogs polling) come from the shared Gateway, DI'd in (composition, not inheritance). -This client adds only the spend surface: /spend/calculate, /spend/tags, -key-spend polling, and the route probes the breadth test uses. +This client adds only the spend surface: /spend/calculate, key-spend +polling, and the route probes the breadth test uses. Re-exports unwrap / is_ok / unique_marker / SpendLogRow so the tests import their helpers from one place. @@ -18,16 +18,10 @@ from dataclasses import dataclass from e2e_config import unique_marker from e2e_http import ( - NetworkError, NoBody, ProbeResult, - RateLimitedError, Result, StreamingResponse, - Success, - UnauthorizedError, - UnknownApiError, - ValidationError, is_ok, unwrap, ) @@ -44,14 +38,10 @@ from models import ( SpendCalculateBody, SpendCalculateResponse, SpendLogRow, - TagSpend, - TagSpends, ) __all__ = [ "SpendClient", - "SpendTagsError", - "TagSpendPoll", "build_client", "reset_spend_logs", "unique_marker", @@ -62,43 +52,6 @@ __all__ = [ ] -class SpendTagsError(Exception): - """A /spend/tags call that did not return 200. Carries the real HTTP status - and body so a server error surfaces verbatim instead of as a missing tag.""" - - def __init__(self, status: int, body: str) -> None: - self.status = status - self.body = body - super().__init__(f"/spend/tags returned HTTP {status}: {body[:500]}") - - -def _spend_tags_error(result: Result[TagSpends]) -> SpendTagsError: - match result: - case Success(): - raise AssertionError("not an error result") - case UnauthorizedError(): - return SpendTagsError(401, "unauthorized") - case RateLimitedError(retry_after_seconds=retry, body=body): - return SpendTagsError(429, f"rate_limited (retry_after={retry}): {body}") - case UnknownApiError(status_code=status, body=body): - return SpendTagsError(status, body) - case NetworkError(message=message): - return SpendTagsError(-1, f"network error: {message}") - case ValidationError(message=message): - return SpendTagsError(-2, f"response did not match TagSpends: {message}") - - -@dataclass(frozen=True, slots=True) -class TagSpendPoll: - """Outcome of polling /spend/tags for one tag: the last matching entry (if any), - whether the endpoint ever returned 200, and a human-readable description of the - last observed HTTP outcome for failure messages.""" - - entry: TagSpend | None - saw_ok: bool - last_status: str - - def reset_spend_logs() -> None: """Truncate LiteLLM_SpendLogs for a clean slate. No proxy endpoint deletes spend logs (/global/spend/reset keeps them), so go to the DB directly. Uses @@ -186,39 +139,6 @@ class SpendClient: ) ).cost - def spend_by_tags(self) -> list[TagSpend]: - """Tagged aggregates from /spend/tags. A non-200 is a real server failure, - not "no tags": raise SpendTagsError with the actual status and body so it - surfaces instead of being swallowed into an empty list.""" - result = self.gateway.transport.get( - "/spend/tags", - headers=self.gateway.transport.master, - params=NoBody(), - response_type=TagSpends, - ) - match result: - case Success(data=data): - return list(data.root) - case _: - raise _spend_tags_error(result) - - def poll_tag_spend(self, tag: str, *, minimum: float = 0.0) -> TagSpendPoll: - """Poll /spend/tags until the tag's aggregate reaches `minimum`. A hard - server error (any non-200) propagates immediately as SpendTagsError rather - than being polled into a timeout; eventual consistency only manifests as a - 200 whose payload does not yet carry the tag, so only that case waits.""" - deadline = time.monotonic() + self.gateway.poll_timeout - entry: TagSpend | None = None - while time.monotonic() < deadline: - tags = self.spend_by_tags() - matches = [t for t in tags if t.individual_request_tag == tag] - if matches: - entry = matches[0] - if (entry.total_spend or 0.0) >= minimum: - return TagSpendPoll(entry=entry, saw_ok=True, last_status="HTTP 200") - time.sleep(self.gateway.poll_interval) - return TagSpendPoll(entry=entry, saw_ok=True, last_status="HTTP 200") - def poll_key_spend(self, key: str, *, minimum: float = 0.0) -> float: deadline = time.monotonic() + self.gateway.poll_timeout spend = 0.0 diff --git a/tests/e2e/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/spend_tracking/test_spend_tracking_e2e.py index 286267a6434..9bd76d1e624 100644 --- a/tests/e2e/spend_tracking/test_spend_tracking_e2e.py +++ b/tests/e2e/spend_tracking/test_spend_tracking_e2e.py @@ -84,18 +84,23 @@ def test_chat_completion_writes_nonzero_spend_row( assert total == prompt + completion, f"token arithmetic broken: {_summarize(rows)}" if chat.id: - assert any(r.request_id == chat.id for r in rows), ( - f"row request_id != client response.id ({chat.id})" - ) + assert any( + r.request_id == chat.id for r in rows + ), f"row request_id != client response.id ({chat.id})" def test_streaming_chat_completion_tracks_spend( client: SpendClient, scoped_key: str ) -> None: result = client.chat_stream( - scoped_key, "gemini-2.5-flash", f"count to three {unique_marker()}", max_tokens=64 + scoped_key, + "gemini-2.5-flash", + f"count to three {unique_marker()}", + max_tokens=64, ) - assert result.ok, f"stream failed (status {result.status_code}): {result.body[:300]}" + assert ( + result.ok + ), f"stream failed (status {result.status_code}): {result.body[:300]}" rows = client.poll_logs_for_key( scoped_key, predicate=lambda rs: any((r.spend or 0) > 0 for r in rs) @@ -105,7 +110,9 @@ def test_streaming_chat_completion_tracks_spend( ) prompt = row.prompt_tokens or 0 completion = row.completion_tokens or 0 - assert prompt > 0 and completion > 0, f"streaming tokens not tracked: {_summarize(rows)}" + assert ( + prompt > 0 and completion > 0 + ), f"streaming tokens not tracked: {_summarize(rows)}" assert (row.total_tokens or 0) == prompt + completion @@ -153,24 +160,27 @@ def test_cache_hit_is_zero_cost_and_suffixed( ) cache_row = cache_rows[0] - assert (cache_row.spend or 0) == 0.0, ( - f"cache hit was charged (double-charge regression): {_summarize(rows)}" - ) + assert ( + cache_row.spend or 0 + ) == 0.0, f"cache hit was charged (double-charge regression): {_summarize(rows)}" assert "_cache_hit" in (cache_row.request_id or ""), ( "cache-hit row missing the _cache_hit request_id suffix; " "duplicate-key collisions will silently drop rows" ) paid_rows = [r for r in rows if r.cache_hit != "True"] - assert any((r.spend or 0) > 0 for r in paid_rows), ( - f"the non-cached call should still be charged: {_summarize(rows)}" - ) + assert any( + (r.spend or 0) > 0 for r in paid_rows + ), f"the non-cached call should still be charged: {_summarize(rows)}" def test_key_spend_equals_sum_of_logs(client: SpendClient, scoped_key: str) -> None: for _ in range(2): _ = unwrap( client.chat( - scoped_key, "gemini-2.5-flash", f"say hi {unique_marker()}", max_tokens=16 + scoped_key, + "gemini-2.5-flash", + f"say hi {unique_marker()}", + max_tokens=16, ) ) @@ -184,15 +194,17 @@ def test_key_spend_equals_sum_of_logs(client: SpendClient, scoped_key: str) -> N assert logs_total > 0 key_spend = client.poll_key_spend(scoped_key, minimum=logs_total * 0.999) - assert _approx_equal(key_spend, logs_total), ( - f"key aggregate {key_spend} != sum of logs {logs_total}; rows: {_summarize(rows)}" - ) + assert _approx_equal( + key_spend, logs_total + ), f"key aggregate {key_spend} != sum of logs {logs_total}; rows: {_summarize(rows)}" def test_request_tags_round_trip(client: SpendClient, scoped_key: str) -> None: tag = f"e2e-spend-{unique_marker()}" _ = unwrap( - client.chat(scoped_key, "gemini-2.5-flash", "tagged request", tags=[tag], max_tokens=16) + client.chat( + scoped_key, "gemini-2.5-flash", "tagged request", tags=[tag], max_tokens=16 + ) ) rows = client.poll_logs_for_key( @@ -203,49 +215,6 @@ def test_request_tags_round_trip(client: SpendClient, scoped_key: str) -> None: ) -def test_tag_spend_matches_sum_of_tagged_logs( - client: SpendClient, scoped_key: str -) -> None: - # Unique tag so /spend/tags can't be polluted by other rows; unique content - # per call so both are fresh misses (paid), not cache hits. - tag = f"e2e-tagspend-{unique_marker()}" - for _ in range(2): - _ = unwrap( - client.chat( - scoped_key, "gemini-2.5-flash", f"hi {unique_marker()}", tags=[tag], max_tokens=16 - ) - ) - - rows = client.poll_logs_for_key( - scoped_key, - min_rows=2, - predicate=lambda rs: sum((r.spend or 0) for r in rs) > 0, - ) - tagged = [r for r in rows if tag in (r.request_tags or [])] - assert len(tagged) >= 2, f"expected 2 tagged rows, saw {_summarize(rows)}" - logs_total = sum((r.spend or 0) for r in tagged) - assert logs_total > 0 - - # A non-200 from /spend/tags raises SpendTagsError (status + body) here, failing - # the test loudly with the real server response instead of masking it as "no - # tags". Reaching this line means the endpoint returned 200 at least once. - poll = client.poll_tag_spend(tag, minimum=logs_total * 0.999) - assert poll.saw_ok, ( - f"/spend/tags never returned 200; last observed: {poll.last_status}" - ) - entry = poll.entry - assert entry is not None, ( - f"tag {tag!r} never appeared in /spend/tags (endpoint was healthy: " - f"{poll.last_status}); tagged rows: {_summarize(tagged)}" - ) - assert _approx_equal(entry.total_spend or 0, logs_total), ( - f"/spend/tags total_spend {entry} != sum of tagged rows {logs_total}" - ) - assert (entry.log_count or 0) == len(tagged), ( - f"/spend/tags log_count {entry.log_count} != tagged rows {len(tagged)}" - ) - - def test_end_user_spend_attributed_on_row( client: SpendClient, scoped_key: str, resources: ResourceManager ) -> None: @@ -283,7 +252,9 @@ def test_failure_call_writes_failure_status_row( def test_spend_calculate_returns_nonzero_cost(client: SpendClient) -> None: - cost = client.calculate_spend("gemini-2.5-flash", "estimate the cost of this request") + cost = client.calculate_spend( + "gemini-2.5-flash", "estimate the cost of this request" + ) assert cost > 0, ( "/spend/calculate returned 0 for gemini-2.5-flash; " "cost map may be missing this model"