diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 54095d85262..e90969da6ea 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -179,10 +179,6 @@ class TagSpends(RootModel[list[TagSpend]]): pass -class SpendTagsResponse(BaseModel): - spend_per_tag: list[TagSpend] | None = None - - # ---------- route probing ---------- diff --git a/tests/e2e/spend_tracking/spend_e2e_client.py b/tests/e2e/spend_tracking/spend_e2e_client.py index d17a9914934..1645297cad4 100644 --- a/tests/e2e/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/spend_tracking/spend_e2e_client.py @@ -17,7 +17,20 @@ from collections.abc import Callable from dataclasses import dataclass from e2e_config import unique_marker -from e2e_http import NoBody, ProbeResult, Result, StreamingResponse, Success, is_ok, unwrap +from e2e_http import ( + NetworkError, + NoBody, + ProbeResult, + RateLimitedError, + Result, + StreamingResponse, + Success, + UnauthorizedError, + UnknownApiError, + ValidationError, + is_ok, + unwrap, +) from e2e_gateway import Gateway, build_gateway from models import ( ChatBody, @@ -31,12 +44,14 @@ from models import ( SpendCalculateBody, SpendCalculateResponse, SpendLogRow, - SpendTagsResponse, TagSpend, + TagSpends, ) __all__ = [ "SpendClient", + "SpendTagsError", + "TagSpendPoll", "build_client", "reset_spend_logs", "unique_marker", @@ -47,6 +62,43 @@ __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 @@ -135,32 +187,37 @@ 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=SpendTagsResponse, + response_type=TagSpends, ) match result: case Success(data=data): - return data.spend_per_tag or [] + return list(data.root) case _: - return [] + raise _spend_tags_error(result) - def poll_tag_spend(self, tag: str, *, minimum: float = 0.0) -> TagSpend | None: - """Poll /spend/tags until the tag's aggregate reaches `minimum`; last seen.""" + 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: - matches = [ - t for t in self.spend_by_tags() if t.individual_request_tag == tag - ] + 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 entry + return TagSpendPoll(entry=entry, saw_ok=True, last_status="HTTP 200") time.sleep(self.gateway.poll_interval) - return entry + 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 diff --git a/tests/e2e/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/spend_tracking/test_spend_tracking_e2e.py index 11f3d8b7591..286267a6434 100644 --- a/tests/e2e/spend_tracking/test_spend_tracking_e2e.py +++ b/tests/e2e/spend_tracking/test_spend_tracking_e2e.py @@ -226,8 +226,18 @@ def test_tag_spend_matches_sum_of_tagged_logs( logs_total = sum((r.spend or 0) for r in tagged) assert logs_total > 0 - entry = client.poll_tag_spend(tag, minimum=logs_total * 0.999) - assert entry is not None, f"tag {tag!r} never appeared in /spend/tags" + # 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}" )