test: update the spend tests that still expected the pre-ledger accounting queues (#50814)

* test: update the spend tests that still expected the pre-ledger accounting queues

#49872 made the SpendLogs ledger the default accounting path on purpose. A batch with no billable request now writes no cost row, a call the guardrail fails closed during a vendor outage keeps one spend row of either status while its key is still charged once, and charges left at a graceful shutdown are applied by the next start. The daily spend batch shutdown tests exercise the old queue, so they pin use_spend_logs_ledger to false.

* test(integration): price the guardrail relay model so the outage test can assert the charge

The relay model had no price, so every stream cost $0 and the outage test could not tell a charged call from an uncharged one. Both rigs now price it, and the test expects each served stream to cost 9 input and 4 output tokens and the key to be charged for all sixteen.

Also notes the skipped batch cost row case in the e2e rules file, next to the other explicit skips.

* test: repair four tests that fail for reasons outside the product

The OpenAPI docs gate test rebuilds the Rust extension with uv, and the integration lane blocks outbound traffic, so uv could not refresh its PyPI index for maturin. The test now runs uv offline against the cache the dependency install already filled.

The response cache test expected a chat reply blocked with empty content to be cached. Since #50057 a reply with no output is never cached, so that case now expects the repeat to reach the upstream and the real answer to be cached.

The Gemini cache hit e2e test gave the model 16 tokens, which its thinking can use up, and a reply with no text is not cached. It now allows 512 and checks the first reply has text.

The Together cache read e2e test picks the cheapest model with a cache read price. Since #50491 that is a -latest alias, whose target reports no cached tokens. Aliases are now left out of the pick.
This commit is contained in:
yuneng-jiang 2026-11-09 22:07:42 -08:00 • committed by GitHub
parent 9a7f5eeec4
commit 07be67603c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 227 additions and 112 deletions

View file

@ -37,7 +37,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
- `embeddings/` - the `/embeddings` endpoint across providers
- `batches/` - the `/batches` endpoint. GovCloud batch coverage is explicitly skipped until its infrastructure is provisioned. Other missing credentials and provider errors remain failures
- `realtime/` - realtime websocket sessions, including the pipecat audio path
- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`)
- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`). The batch cost row attribution case in `spend_tracking/` is explicitly skipped, because the SpendLogs ledger writes no cost row for a $0 batch and a batch with a cost does not finish within one run
- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface, not Playwright)
- `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0)
- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below). The semantic tool filter test turns on a proxy-wide setting, so it is marked `proxy_wide_config` and runs on the redeployed managed-files stack after the main run, never next to it

View file

@ -1617,17 +1617,6 @@ class TestBatchFailurePaths:
assert first_error.message, "batch error item has no message"
assert first_error.code, "batch error item has no code"
rows = client.proxy.poll_logs_for_request_id(f"{fetched.id}_batch_cost")
assert rows, (
f"failed batch {fetched.id} wrote no spend row; retrieve must book $0 (LIT-4852)"
)
assert all((row.spend or 0) == 0 for row in rows), (
f"failed batch must cost $0, got {[(r.request_id, r.spend) for r in rows]}"
)
assert rows[0].call_type == "aretrieve_batch", (
f"batch cost row call_type={rows[0].call_type!r}"
)
conflict = client.cancel_batch(batch.id, key=key)
match conflict:
case UnknownApiError(status_code=409, body=body):

View file

@ -166,6 +166,7 @@ def _cheapest_together_chat_model(registry: Mapping[str, CostMapEntry], needs: _
def qualifies(name: str, entry: CostMapEntry) -> bool:
return (
name.startswith("together_ai/")
and not name.endswith("-latest")
and entry.litellm_provider == "together_ai"
and entry.mode == "chat"
and (entry.deprecation_date is None or entry.deprecation_date > today)

View file

@ -410,6 +410,9 @@ class TestKeyAttribution:
rehashed: Final = [(row.request_id, row.api_key) for row in rows if row.api_key != HEALTH_SERVICE_ACCOUNT]
assert not rehashed, f"health-check rows keyed by something other than {HEALTH_SERVICE_ACCOUNT!r}: {rehashed}"
@pytest.mark.skip(
reason="the SpendLogs ledger writes no cost row for a $0 batch, and a batch with a cost is out of one run's reach"
)
@pytest.mark.covers(
"quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key",
exercised_on=["batches"],

View file

@ -287,8 +287,12 @@ def test_cache_hit_is_zero_cost_and_suffixed(
# populated. The marker keeps each run isolated - a fixed prompt would persist
# in the shared response cache across runs and make both calls hit (flaky).
prompt = f"What is the capital of France? Answer in one word. {unique_marker()}"
_ = unwrap(client.chat(scoped_key, GEMINI_MODEL, prompt, max_tokens=16, cache=None))
_ = unwrap(client.chat(scoped_key, GEMINI_MODEL, prompt, max_tokens=16, cache=None))
first = unwrap(client.chat(scoped_key, GEMINI_MODEL, prompt, max_tokens=512, cache=None))
answered = first.choices and first.choices[0].message and first.choices[0].message.content
assert answered, (
f"{GEMINI_MODEL} returned no text, and a reply with no output is never cached: {first}"
)
_ = unwrap(client.chat(scoped_key, GEMINI_MODEL, prompt, max_tokens=512, cache=None))
rows = client.poll_logs_for_key(
scoped_key,

View file

@ -603,34 +603,45 @@ def test_embeddings_are_still_served_from_the_cache(gateway: Gateway) -> None:
assert second is not None and second["data"] == first["data"], second
@pytest.mark.parametrize(
"choice",
(
pytest.param(
{"index": 0, "message": {"role": "assistant", "content": ""}, "finish_reason": "content_filter"},
id="blocked-empty-content",
),
pytest.param(
{
"index": 0,
"message": {
"role": "assistant",
"content": None,
"tool_calls": [
{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}
],
},
"finish_reason": "tool_calls",
},
id="tool-calls-only",
),
),
)
def test_chat_answers_with_a_choice_but_no_text_are_still_cached(
gateway: Gateway, choice: dict[str, JsonValue]
def test_chat_answer_blocked_with_empty_content_reaches_upstream_again_and_the_real_answer_is_cached(
gateway: Gateway,
) -> None:
upstream: Final = Upstream(gateway.upstream_url)
store: Final = redis_store()
blocked: Final[dict[str, JsonValue]] = {
"index": 0,
"message": {"role": "assistant", "content": ""},
"finish_reason": "content_filter",
}
with gateway.scenario() as scenario:
handle: Final = scripted(scenario, json_response({**CHAT, "choices": [blocked]}))
model: Final = scenario.model(api_base=handle.api_base())
request: Final = chat_body(model, prompt())
first: Final = _chat(gateway, request)
assert object_value(choices(first)[0])["finish_reason"] == "content_filter", first.text
assert upstream.calls(handle.scenario_id) == 1
rescript(handle, json_response(CHAT))
real: Final = _chat(gateway, request)
assert response_id(real) != response_id(first) and content_of_first_choice(real) == "scripted", real.text
assert upstream.calls(handle.scenario_id) == 2
await_entry(store, response_id(real))
hit: Final = _chat(gateway, request)
assert response_id(hit) == response_id(real) and content_of_first_choice(hit) == "scripted", hit.text
assert upstream.calls(handle.scenario_id) == 2
def test_chat_answers_with_tool_calls_but_no_text_are_still_cached(gateway: Gateway) -> None:
upstream: Final = Upstream(gateway.upstream_url)
store: Final = redis_store()
choice: Final[dict[str, JsonValue]] = {
"index": 0,
"message": {
"role": "assistant",
"content": None,
"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}],
},
"finish_reason": "tool_calls",
}
with gateway.scenario() as scenario:
handle: Final = scripted(scenario, json_response({**CHAT, "choices": [choice]}))
model: Final = scenario.model(api_base=handle.api_base())

View file

@ -1,9 +1,11 @@
from typing import Final
import openapi_docs_gate as gate
import pytest
def test_base_counts_imports_native_extension_from_detached_worktree() -> None:
def test_base_counts_imports_native_extension_from_detached_worktree(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("UV_OFFLINE", "1")
expected_counts: Final = dict(gate.count_by_rule(gate.head_violations()))
assert expected_counts
assert dict(gate.base_counts("HEAD")) == expected_counts

View file

@ -7,8 +7,10 @@ import time
from hashlib import sha256
from typing import Final
import httpx
import pytest
import websockets
from _pytest.mark.structures import ParameterSet
from integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value, string_value
from integration._support.upstream import delete_scenario, register_scenario
from integration.cost_calculation.assertions import assert_exact
@ -150,46 +152,54 @@ def _batch_input_lines(case: BatchCostCase, model_name: str) -> bytes:
).encode()
@pytest.mark.parametrize(
"case",
tuple(pytest.param(case, marks=pytest.mark.covers(case.covers), id=case.name) for case in BATCH_CASES),
)
def _batch_retrieved_twice(
gateway: Gateway, scenario: Scenario, case: BatchCostCase
) -> tuple[str, str, httpx.Response]:
key: Final = scenario.key()
model_name, identity = _register_deployment(
scenario,
case.litellm_model,
_batch_response(case),
case.name,
realtime=False,
)
file_response: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "model": model_name},
{"file": ("in.jsonl", _batch_input_lines(case, model_name), "application/jsonl")},
key=key,
)
assert file_response.is_success, file_response.text
file_body: Final = JSON_OBJECT.validate_json(file_response.content)
batch_response: Final = gateway.request(
"POST",
"/v1/batches",
{
"input_file_id": string_value(file_body["id"]),
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"model": model_name,
},
key=key,
)
assert batch_response.is_success, batch_response.text
batch_body: Final = JSON_OBJECT.validate_json(batch_response.content)
batch_id: Final = string_value(batch_body["id"])
first_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key)
second_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key)
assert first_retrieval.is_success, first_retrieval.text
assert second_retrieval.is_success, second_retrieval.text
return key, identity, second_retrieval
def _batch_params(cases: tuple[BatchCostCase, ...]) -> tuple[ParameterSet, ...]:
return tuple(pytest.param(case, marks=pytest.mark.covers(case.covers), id=case.name) for case in cases)
@pytest.mark.parametrize("case", _batch_params(tuple(case for case in BATCH_CASES if case.expected.spend > 0)))
def test_batch_costs(gateway: Gateway, case: BatchCostCase) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
model_name, identity = _register_deployment(
scenario,
case.litellm_model,
_batch_response(case),
case.name,
realtime=False,
)
file_response: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "model": model_name},
{"file": ("in.jsonl", _batch_input_lines(case, model_name), "application/jsonl")},
key=key,
)
assert file_response.is_success, file_response.text
file_body: Final = JSON_OBJECT.validate_json(file_response.content)
batch_response: Final = gateway.request(
"POST",
"/v1/batches",
{
"input_file_id": string_value(file_body["id"]),
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"model": model_name,
},
key=key,
)
assert batch_response.is_success, batch_response.text
batch_body: Final = JSON_OBJECT.validate_json(batch_response.content)
batch_id: Final = string_value(batch_body["id"])
first_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key)
second_retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key)
assert first_retrieval.is_success, first_retrieval.text
assert second_retrieval.is_success, second_retrieval.text
key, identity, second_retrieval = _batch_retrieved_twice(gateway, scenario, case)
retrieval_rows: Final = poll_rows_where(key, 1, lambda row: row.call_type == "aretrieve_batch")
assert len(retrieval_rows) == 1
rows: Final = read_rows_now(key)
@ -203,6 +213,22 @@ def test_batch_costs(gateway: Gateway, case: BatchCostCase) -> None:
assert len(tuple(row for row in read_rows_now(key) if row.call_type == "aretrieve_batch")) == 1
@pytest.mark.parametrize("case", _batch_params(tuple(case for case in BATCH_CASES if case.expected.spend == 0)))
def test_batch_with_no_billable_request_is_retrieved_free_and_writes_no_cost_row(
gateway: Gateway, case: BatchCostCase
) -> None:
with gateway.scenario() as scenario:
key, _, _ = _batch_retrieved_twice(gateway, scenario, case)
later_model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
later: Final = gateway.chat(later_model, key=key, text="sent after the batch retrievals")
assert later["usage"] == {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40}, later
rows: Final = poll_rows_where(key, 1, lambda row: row.call_type == "acompletion")
assert [row.spend for row in rows] == [pytest.approx(20 * 0.001 + 20 * 0.002)], rows
every_row: Final = read_rows_now(key)
assert [row.call_type for row in every_row if row.call_type == "aretrieve_batch"] == [], every_row
assert sum(row.spend or 0.0 for row in every_row) == pytest.approx(20 * 0.001 + 20 * 0.002), every_row
def _realtime_response(case: RealtimeCostCase) -> RealtimeResponse:
return RealtimeResponse(
content_type="application/x-realtime",

View file

@ -11,6 +11,7 @@ request as before."""
from __future__ import annotations
import hashlib
import json
import os
import threading
@ -39,6 +40,8 @@ from integration._support.wire import Reply, Request, Wire, wire_server
pytestmark: Final = pytest.mark.shares_worker_state
MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0"
INPUT_COST_PER_TOKEN: Final = 0.001
OUTPUT_COST_PER_TOKEN: Final = 0.002
ROUTER_MODEL: Final = "converse-relay"
EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
SAMPLED: Final = "relay-sampled"
@ -390,6 +393,10 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
"aws_secret_access_key": "synthetic-relay-secret",
"aws_region_name": "us-east-1",
},
"model_info": {
"input_cost_per_token": INPUT_COST_PER_TOKEN,
"output_cost_per_token": OUTPUT_COST_PER_TOKEN,
},
}
]
config_path: Final = directory / "converse-stream-relay.yaml"
@ -651,6 +658,24 @@ def _spend_statuses(call_id: str, at_least: int) -> list[JsonValue]:
return [row["status"] for row in rows]
def _logged_spend(call_id: str) -> float:
rows: Final = read_rows(
'SELECT spend FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s OR request_id=%s '
"OR metadata->>'litellm_call_id'=%s",
(call_id, call_id, call_id),
)
assert len(rows) == 1 and isinstance(rows[0]["spend"], float), rows
return rows[0]["spend"]
def _key_spend(key: str) -> float:
rows: Final = read_rows(
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashlib.sha256(key.encode()).hexdigest(),)
)
assert len(rows) == 1 and isinstance(rows[0]["spend"], float), rows
return rows[0]["spend"]
def _exception(frames: Sequence[Frame]) -> tuple[str, str] | None:
last: Final = frames[-1] if frames else None
if last is None or last.message_type != "exception":
@ -1001,6 +1026,7 @@ def test_guardrail_vendor_outage_during_a_burst_fails_closed_then_recovers(rig:
with ThreadPoolExecutor(max_workers=count) as pool:
return list(zip(markers, pool.map(partial(_stream_with_call_id, rig, OUTAGE), markers)))
spend_before: Final = _key_spend(rig.keys[OUTAGE])
rig.outage.set()
try:
during: Final = burst(8)
@ -1013,5 +1039,17 @@ def test_guardrail_vendor_outage_during_a_burst_fails_closed_then_recovers(rig:
for marker, (status, _, body, _) in after:
assert status == 200 and _exception(decode(body)) is None, (marker, body)
assert body == b"".join(rig.scripts[marker].chunks or ()), (marker, decode(body))
for marker, (_, _, _, call_id) in (*during, *after):
assert _spend_statuses(call_id, 1).count("success") == 1, marker
for marker, (_, _, _, call_id) in after:
assert _spend_statuses(call_id, 1) == ["success"], marker
for marker, (_, _, _, call_id) in during:
assert len(_spend_statuses(call_id, 1)) == 1, marker
stream_cost: Final = 9 * INPUT_COST_PER_TOKEN + 4 * OUTPUT_COST_PER_TOKEN
costs: Final = [_logged_spend(call_id) for _, (_, _, _, call_id) in after]
assert costs == [pytest.approx(stream_cost)] * len(after), costs
charged: Final = eventually(
lambda: _key_spend(rig.keys[OUTAGE]) - spend_before,
lambda spend: spend == pytest.approx(len((*during, *after)) * stream_cost),
seconds=70,
return_last_on_timeout=True,
)
assert charged == pytest.approx(len((*during, *after)) * stream_cost), charged

View file

@ -12,6 +12,7 @@ Streams of other model families keep going out as received"""
from __future__ import annotations
import base64
import hashlib
import json
import os
import threading
@ -43,6 +44,8 @@ pytestmark: Final = pytest.mark.shares_worker_state
CLAUDE_ID: Final = "anthropic.claude-sonnet-5-v1:0"
LLAMA_ID: Final = "meta.llama4-maverick-17b-instruct-v1:0"
PROFILE_ID: Final = "arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/relayprofile"
INPUT_COST_PER_TOKEN: Final = 0.001
OUTPUT_COST_PER_TOKEN: Final = 0.002
CLAUDE: Final = "invoke-relay"
LLAMA: Final = "invoke-relay-llama"
PROFILE: Final = "invoke-relay-profile"
@ -334,6 +337,7 @@ def _deployment(name: str, model_id: str, upstream: Wire) -> dict[str, JsonValue
"aws_secret_access_key": "synthetic-relay-secret",
"aws_region_name": "us-east-1",
},
"model_info": {"input_cost_per_token": INPUT_COST_PER_TOKEN, "output_cost_per_token": OUTPUT_COST_PER_TOKEN},
}
@ -468,6 +472,24 @@ def _stream(rig: Rig, guardrail: str, marker: str, model: str = CLAUDE) -> tuple
return status, content_type, body
def _logged_spend(call_id: str) -> float:
rows: Final = read_rows(
'SELECT spend FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s OR request_id=%s '
"OR metadata->>'litellm_call_id'=%s",
(call_id, call_id, call_id),
)
assert len(rows) == 1 and isinstance(rows[0]["spend"], float), rows
return rows[0]["spend"]
def _key_spend(key: str) -> float:
rows: Final = read_rows(
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashlib.sha256(key.encode()).hexdigest(),)
)
assert len(rows) == 1 and isinstance(rows[0]["spend"], float), rows
return rows[0]["spend"]
def _exception(frames: Sequence[Frame]) -> tuple[str, str] | None:
last: Final = frames[-1] if frames else None
if last is None or last.message_type != "exception":
@ -807,6 +829,7 @@ def test_guardrail_vendor_outage_during_a_burst_fails_closed_then_recovers(rig:
with ThreadPoolExecutor(max_workers=count) as pool:
return list(zip(markers, pool.map(partial(_stream_with_call_id, rig, OUTAGE), markers)))
spend_before: Final = _key_spend(rig.keys[OUTAGE])
rig.outage.set()
try:
during: Final = burst(8)
@ -819,8 +842,20 @@ def test_guardrail_vendor_outage_during_a_burst_fails_closed_then_recovers(rig:
for marker, (status, _, body, _) in after:
assert status == 200 and _exception(decode(body)) is None, (marker, body)
assert body == b"".join(rig.scripts[marker].chunks or ()), (marker, decode(body))
for marker, (_, _, _, call_id) in (*during, *after):
assert _spend_statuses(call_id, 1).count("success") == 1, marker
for marker, (_, _, _, call_id) in after:
assert _spend_statuses(call_id, 1) == ["success"], marker
for marker, (_, _, _, call_id) in during:
assert len(_spend_statuses(call_id, 1)) == 1, marker
stream_cost: Final = 9 * INPUT_COST_PER_TOKEN + 4 * OUTPUT_COST_PER_TOKEN
costs: Final = [_logged_spend(call_id) for _, (_, _, _, call_id) in after]
assert costs == [pytest.approx(stream_cost)] * len(after), costs
charged: Final = eventually(
lambda: _key_spend(rig.keys[OUTAGE]) - spend_before,
lambda spend: spend == pytest.approx(len((*during, *after)) * stream_cost),
seconds=70,
return_last_on_timeout=True,
)
assert charged == pytest.approx(len((*during, *after)) * stream_cost), charged
def test_client_disconnect_stops_the_proxy_reading_the_upstream(rig: Rig) -> None:

View file

@ -94,14 +94,19 @@ def _tagged_chat(gateway: Gateway, key: str, tags: tuple[str, str]) -> None:
assert response.status_code == 200, response.text
def _assert_requests_and_spend(table: str, count: str, key: str, requests: int) -> None:
hashed: Final = hashlib.sha256(key.encode()).hexdigest()
rows: Final = read_rows(
f'SELECT ({count})::int AS requests, sum(spend) AS spend FROM "{table}" WHERE api_key=%s', (hashed,)
)
assert rows[0]["requests"] == requests, (table, rows)
assert rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (table, rows)
def _assert_shutdown_accounting(key: str, tags: tuple[str, str], requests: int) -> None:
hashed: Final = hashlib.sha256(key.encode()).hexdigest()
for table, count in (("LiteLLM_SpendLogs", "count(*)"), ("LiteLLM_DailyUserSpend", "sum(api_requests)")):
rows: Final = read_rows(
f'SELECT ({count})::int AS requests, sum(spend) AS spend FROM "{table}" WHERE api_key=%s', (hashed,)
)
assert rows[0]["requests"] == requests, (table, rows)
assert rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (table, rows)
_assert_requests_and_spend("LiteLLM_SpendLogs", "count(*)", key, requests)
_assert_requests_and_spend("LiteLLM_DailyUserSpend", "sum(api_requests)", key, requests)
key_rows: Final = read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashed,))
assert key_rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9)
for tag in tags:
@ -137,7 +142,7 @@ def test_tagged_request_survives_shutdown_before_the_tag_timer(
@pytest.mark.parametrize("redis_buffer", [False, True])
@pytest.mark.parametrize("overflow", [False, True])
def test_callback_backlog_is_accounted_before_shutdown_disconnects(
def test_callback_backlog_is_logged_before_shutdown_disconnects_and_accounted_once_by_the_next_start(
gateway: Gateway, tmp_path: Path, redis_buffer: bool, overflow: bool
) -> None:
release: Final = threading.Event()
@ -148,19 +153,15 @@ def test_callback_backlog_is_accounted_before_shutdown_disconnects(
wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release)) as logger,
):
config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer, gated=True)
with owned_proxy_process(
gateway,
tmp_path,
{
"DATABASE_URL": os.environ["DATABASE_URL"],
"LOGGING_WORKER_CONCURRENCY": "2",
"LOGGING_WORKER_MAX_QUEUE_SIZE": "2" if overflow else "50000",
"SHUTDOWN_CALLBACK_URL": logger.url,
"GRACEFUL_SHUTDOWN_TIMEOUT": "15",
"LITELLM_LOG": "DEBUG",
},
config=config,
) as owned:
overrides: Final = {
"DATABASE_URL": os.environ["DATABASE_URL"],
"LOGGING_WORKER_CONCURRENCY": "2",
"LOGGING_WORKER_MAX_QUEUE_SIZE": "2" if overflow else "50000",
"SHUTDOWN_CALLBACK_URL": logger.url,
"GRACEFUL_SHUTDOWN_TIMEOUT": "15",
"LITELLM_LOG": "DEBUG",
}
with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned:
key: Final = string_value(
owned.gateway.post("/key/generate", {"user_id": f"backlog-{uuid.uuid4().hex}"})["key"]
)
@ -179,9 +180,13 @@ def test_callback_backlog_is_accounted_before_shutdown_disconnects(
)
release.set()
owned.process.wait(timeout=90)
_assert_shutdown_accounting(key, tags, requests)
finally:
release.set()
_assert_requests_and_spend("LiteLLM_SpendLogs", "count(*)", key, requests)
with owned_proxy_process(gateway, tmp_path, overrides, config=config) as restarted:
restarted.begin_shutdown()
restarted.process.wait(timeout=90)
_assert_shutdown_accounting(key, tags, requests)
@pytest.mark.parametrize("gunicorn", [False, True])
@ -396,17 +401,18 @@ class _Shutdown:
self.owned.process.wait(timeout=120)
def _config_with_pool_limit(tmp_path: Path, pool_limit: int) -> Path:
def _legacy_spend_queue_config_with_pool_limit(tmp_path: Path, pool_limit: int) -> Path:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["general_settings"]["database_connection_pool_limit"] = pool_limit
config["general_settings"]["database_connection_pool_timeout"] = 60
config["general_settings"]["use_spend_logs_ledger"] = False
path: Final = tmp_path / f"pool-{pool_limit}.yaml"
path.write_text(yaml.safe_dump(config))
return path
@contextmanager
def _proxy_with_one_seeded_row(
def _legacy_spend_queue_proxy_with_one_seeded_row(
gateway: Gateway,
tmp_path: Path,
pool_limit: int,
@ -429,7 +435,7 @@ def _proxy_with_one_seeded_row(
"SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS": "1",
"SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS": str(cancel_timeout_seconds),
},
config=_config_with_pool_limit(tmp_path, pool_limit),
config=_legacy_spend_queue_config_with_pool_limit(tmp_path, pool_limit),
remove_environment=("DATABASE_URL_READ_REPLICA",),
workers=workers,
) as owned:
@ -460,7 +466,7 @@ def test_daily_spend_batch_cancelled_while_waiting_for_a_pool_connection_is_writ
gateway: Gateway, tmp_path: Path
) -> None:
with (
_proxy_with_one_seeded_row(gateway, tmp_path, pool_limit=2) as shutdown,
_legacy_spend_queue_proxy_with_one_seeded_row(gateway, tmp_path, pool_limit=2) as shutdown,
psycopg.connect(os.environ["DATABASE_URL"]) as models,
psycopg.connect(os.environ["DATABASE_URL"]) as memberships,
):
@ -482,7 +488,7 @@ def test_daily_spend_batch_cancelled_while_waiting_for_a_row_lock_is_written_exa
gateway: Gateway, tmp_path: Path
) -> None:
with (
_proxy_with_one_seeded_row(gateway, tmp_path, pool_limit=10) as shutdown,
_legacy_spend_queue_proxy_with_one_seeded_row(gateway, tmp_path, pool_limit=10) as shutdown,
psycopg.connect(os.environ["DATABASE_URL"]) as holder,
psycopg.connect(os.environ["DATABASE_URL"]) as memberships,
):
@ -508,7 +514,7 @@ def test_daily_spend_batch_cancelled_while_postgres_is_committing_it_is_written_
gateway: Gateway, tmp_path: Path, cancel_timeout_seconds: int, commit_fails_once: bool
) -> None:
with (
_proxy_with_one_seeded_row(
_legacy_spend_queue_proxy_with_one_seeded_row(
gateway, tmp_path, pool_limit=10, cancel_timeout_seconds=cancel_timeout_seconds, settle_seconds=90
) as shutdown,
psycopg.connect(os.environ["DATABASE_URL"]) as memberships,
@ -528,7 +534,7 @@ def test_daily_spend_batch_cancelled_while_postgres_is_committing_it_is_written_
def test_daily_spend_burst_across_two_workers_survives_shutdown_during_commit_exactly_once(
gateway: Gateway, tmp_path: Path
) -> None:
with _proxy_with_one_seeded_row(
with _legacy_spend_queue_proxy_with_one_seeded_row(
gateway, tmp_path, pool_limit=10, settle_seconds=120, requests=BURST_REQUESTS, workers=2
) as shutdown:
suffix: Final = _install_slow_commit(shutdown.owner, fails_once=False)

View file

@ -278,7 +278,7 @@ def test_member_added_without_any_budget_is_charged_on_its_membership_row(gatewa
@pytest.mark.covers("spend.team_member.stale_low_redis_counter_still_blocks_member_over_budget")
def test_member_over_budget_is_blocked_when_redis_counter_reads_stale_low(gateway: Gateway) -> None:
def test_member_over_budget_is_blocked_when_redis_counter_reads_stale_low(gateway: Gateway, peer: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
@ -308,7 +308,7 @@ def test_member_over_budget_is_blocked_when_redis_counter_reads_stale_low(gatewa
assert float(counted) == pytest.approx(0.06), counted
cache.set(counter_key, "0.01")
upstream.delete("/__observations").raise_for_status()
denied: Final = gateway.request(
denied: Final = peer.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"stale counter {uuid.uuid4().hex}"}]},