From 07be67603c08918140d1fde65874eb098adbc991 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 9 Nov 2026 22:07:42 -0800 Subject: [PATCH] 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. --- tests/e2e/AGENTS.md | 2 +- tests/e2e/batches/test_batches_e2e.py | 11 -- .../llm_translation/test_together_ai_e2e.py | 1 + .../test_key_attribution_e2e.py | 3 + .../spend_tracking/test_spend_tracking_e2e.py | 8 +- .../test_response_cache_empty_output.py | 61 ++++++----- .../configuration/test_openapi_docs_gate.py | 4 +- .../test_batch_realtime_cost.py | 102 +++++++++++------- ...ock_converse_stream_post_call_guardrail.py | 42 +++++++- ...drock_invoke_stream_post_call_guardrail.py | 39 ++++++- .../integration/spend/test_shutdown_flush.py | 62 ++++++----- .../spend/test_team_member_spend.py | 4 +- 12 files changed, 227 insertions(+), 112 deletions(-) diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 89bf24016da..abdc1ecafbb 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -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 diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index eb83db73199..202d28c829a 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -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): diff --git a/tests/e2e/llm_translation/test_together_ai_e2e.py b/tests/e2e/llm_translation/test_together_ai_e2e.py index 7775c5c6ed2..9ccdf7fa2d2 100644 --- a/tests/e2e/llm_translation/test_together_ai_e2e.py +++ b/tests/e2e/llm_translation/test_together_ai_e2e.py @@ -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) diff --git a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py index 1d0878d073a..66c5d31c244 100644 --- a/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_key_attribution_e2e.py @@ -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"], 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 d452525ec91..ea20b4a40b6 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 @@ -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, diff --git a/tests/integration/caching/test_response_cache_empty_output.py b/tests/integration/caching/test_response_cache_empty_output.py index 7a20e712d4e..c7afaea36f1 100644 --- a/tests/integration/caching/test_response_cache_empty_output.py +++ b/tests/integration/caching/test_response_cache_empty_output.py @@ -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()) diff --git a/tests/integration/configuration/test_openapi_docs_gate.py b/tests/integration/configuration/test_openapi_docs_gate.py index eefa663cb10..1afa450c67a 100644 --- a/tests/integration/configuration/test_openapi_docs_gate.py +++ b/tests/integration/configuration/test_openapi_docs_gate.py @@ -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 diff --git a/tests/integration/cost_calculation/test_batch_realtime_cost.py b/tests/integration/cost_calculation/test_batch_realtime_cost.py index 1cff2ffb4ed..0e7de18a9a7 100644 --- a/tests/integration/cost_calculation/test_batch_realtime_cost.py +++ b/tests/integration/cost_calculation/test_batch_realtime_cost.py @@ -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", diff --git a/tests/integration/observability/test_bedrock_converse_stream_post_call_guardrail.py b/tests/integration/observability/test_bedrock_converse_stream_post_call_guardrail.py index 9546b9549ff..9ed42b4633c 100644 --- a/tests/integration/observability/test_bedrock_converse_stream_post_call_guardrail.py +++ b/tests/integration/observability/test_bedrock_converse_stream_post_call_guardrail.py @@ -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 diff --git a/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py b/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py index ae243324e27..f7f5a530cf7 100644 --- a/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py +++ b/tests/integration/observability/test_bedrock_invoke_stream_post_call_guardrail.py @@ -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: diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index 174403d1873..17b835a2d84 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -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) diff --git a/tests/integration/spend/test_team_member_spend.py b/tests/integration/spend/test_team_member_spend.py index c11dc9db4bf..d84ffaf9325 100644 --- a/tests/integration/spend/test_team_member_spend.py +++ b/tests/integration/spend/test_team_member_spend.py @@ -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}"}]},