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}"}]},