mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
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.
This commit is contained in:
mubashir1osmani
parent
3f36a64355
commit
0e84aa5b5a
4 changed files with 44 additions and 163 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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) |
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue