From f925c067d2dc47d60e730965d0df6eeb38fc17d9 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:07:06 +0000 Subject: [PATCH] fix(tests): read owned proxy rows from the owned proxy's database (#47333) * fix(tests): read the stream stall cooldown spend rows from the owned proxy's database The cooldown cells boot an owned proxy, which writes its spend logs to its own scratch database since owned proxies went private, while the cells polled the worker's database and never saw a row. Read each row through the gateway that served the request, so the cells check the proxy they called. * fix(tests): read passthrough spend rows from the owned proxy's database The vertex regional uplift and passthrough stream pricing spend files boot private owned proxies, then polled the worker database for spend rows that only reach the proxy's own database. Read them through the owned proxy's gateway instead. Also hold the stream stall file's hold_headers peer past the longest row wait, so a slow spend row cannot fail the peer teardown. * fix(tests): check agent deletion on the owned proxy's database The agent least privilege test deletes its agent through an owned proxy, then checked the worker database for the row, which is always empty there. Read it through the owned proxy's gateway so a delete that leaves the row fails the test. * fix(tests): read turn admission spend rows from the owned proxy's database --- .../test_agent_least_privilege.py | 3 +- .../test_hosted_vllm_stream_stall_wire.py | 2 +- ...test_responses_websocket_turn_admission.py | 11 +++--- ..._passthrough_stream_pricing_paths_spend.py | 38 ++++++++++--------- ...ertex_passthrough_regional_uplift_spend.py | 9 ++--- 5 files changed, 31 insertions(+), 32 deletions(-) diff --git a/tests/integration/authorization/test_agent_least_privilege.py b/tests/integration/authorization/test_agent_least_privilege.py index 8a7f1596379..498a53c3ca9 100644 --- a/tests/integration/authorization/test_agent_least_privilege.py +++ b/tests/integration/authorization/test_agent_least_privilege.py @@ -8,7 +8,6 @@ import httpx import pytest import yaml from integration._support.client import Gateway, eventually, object_value, string_value -from integration._support.database import read_rows from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server @@ -108,7 +107,7 @@ def _create_agent(candidate: Gateway, wire: Wire, agent_name: str, cleanups: Exi def _delete_agent(candidate: Gateway, agent_id: str) -> None: deleted: Final = candidate.request("DELETE", f"/v1/agents/{agent_id}") assert deleted.status_code == 200, deleted.text - assert read_rows('SELECT agent_id FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (agent_id,)) == [] + assert candidate.rows('SELECT agent_id FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (agent_id,)) == [] def _card(candidate: Gateway, agent_id: str, key: str) -> httpx.Response: diff --git a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py index 00e5388b474..fad7898871d 100644 --- a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py +++ b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py @@ -42,8 +42,8 @@ _RESPONSES_TARGET: Final = "/v1/responses" _STREAM_TIMEOUT_SECONDS: Final = 1 _DEPLOYMENT_TIMEOUT_SECONDS: Final = 30 _PAUSE_SECONDS: Final = 3.0 -_HOLD_SECONDS: Final = 30.0 _ROW_SECONDS: Final = 60 +_HOLD_SECONDS: Final = 2.0 * _ROW_SECONDS _EXTRA: Final[dict[str, JsonValue]] = {"num_retries": 0, "cache": {"no-cache": True}} _ROWS: Final = 'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE model_group=%s' _EITHER_GROUP_ROWS: Final = 'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE model_group IN (%s, %s)' diff --git a/tests/integration/routing/test_responses_websocket_turn_admission.py b/tests/integration/routing/test_responses_websocket_turn_admission.py index a246c89e7b3..87b56a92979 100644 --- a/tests/integration/routing/test_responses_websocket_turn_admission.py +++ b/tests/integration/routing/test_responses_websocket_turn_admission.py @@ -27,7 +27,6 @@ from integration._support.client import ( object_value, string_value, ) -from integration._support.database import read_rows from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.responses_ws_peer import ( PROVIDER_MODEL, @@ -260,9 +259,9 @@ async def _until(read: Callable[[], float], satisfied: Callable[[float], bool]) return await asyncio.to_thread(eventually, read, satisfied, 15) -def _failure_rows(key: str) -> list[dict[str, JsonValue]]: +def _failure_rows(gateway: Gateway, key: str) -> list[dict[str, JsonValue]]: rows: Final = eventually( - lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), + lambda: gateway.rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), lambda found: any(row["status"] == "failure" for row in found), seconds=70, ) @@ -291,7 +290,7 @@ def test_a_turn_past_the_deployment_budget_is_refused_on_the_open_socket_and_log _assert_refused(second) _assert_refused(third) assert deployment.forwarded_turns() == 1 - failures: Final = _failure_rows(key) + failures: Final = _failure_rows(limited, key) assert len(failures) == 1, failures assert failures[0]["model_group"] == deployment.name, failures assert tag in list_value(failures[0]["request_tags"]), failures @@ -355,7 +354,7 @@ def test_a_turn_whose_admission_check_fails_gets_a_500_frame_and_the_socket_serv _assert_turn_check_failed(third, fault) completed(fourth) assert deployment.forwarded_turns() == 2 - failures: Final = _failure_rows(key) + failures: Final = _failure_rows(faulting, key) assert len(failures) == 1, failures assert failures[0]["model_group"] == deployment.name, failures recorded: Final = object_value(object_value(failures[0]["metadata"])["error_information"]) @@ -387,7 +386,7 @@ def test_a_turn_still_streaming_when_the_next_turns_check_fails_completes_and_is billed: Final = completed(streamed) assert deployment.forwarded_turns() == 1 rows: Final = eventually( - lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), + lambda: faulting.rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)), lambda found: any(row["request_id"] == billed["id"] for row in found), seconds=70, ) diff --git a/tests/integration/spend/test_passthrough_stream_pricing_paths_spend.py b/tests/integration/spend/test_passthrough_stream_pricing_paths_spend.py index f091a8f8376..0ead544a2ee 100644 --- a/tests/integration/spend/test_passthrough_stream_pricing_paths_spend.py +++ b/tests/integration/spend/test_passthrough_stream_pricing_paths_spend.py @@ -12,7 +12,6 @@ import httpx import pytest import yaml from integration._support.client import Gateway, eventually, gateway_from_environment -from integration._support.database import read_rows from integration._support.forward_proxy import tunnelling_forward_proxy from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.tls import server_context, write_self_signed_cert @@ -289,9 +288,9 @@ def gateway_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: yield _Rig(proxy, upstream) -def _spend_row(key: str) -> dict[str, JsonValue]: +def _spend_row(gateway: Gateway, key: str) -> dict[str, JsonValue]: rows: Final = eventually( - lambda: read_rows( + lambda: gateway.rows( 'SELECT spend, custom_llm_provider FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),), ), @@ -301,14 +300,14 @@ def _spend_row(key: str) -> dict[str, JsonValue]: return rows[0] -def _spend_for_key(key: str) -> float: - spend: Final = _spend_row(key)["spend"] +def _spend_for_key(gateway: Gateway, key: str) -> float: + spend: Final = _spend_row(gateway, key)["spend"] assert isinstance(spend, float), key return spend -def _provider_for_key(key: str) -> JsonValue: - return _spend_row(key)["custom_llm_provider"] +def _provider_for_key(gateway: Gateway, key: str) -> JsonValue: + return _spend_row(gateway, key)["custom_llm_provider"] def _post_and_read(rig: _Rig, path: str, body: Mapping[str, JsonValue], headers: Mapping[str, str]) -> None: @@ -347,7 +346,7 @@ def test_claude_on_vertex_passthrough_spend_carries_the_regional_uplift_whether_ {**body, "stream": action == "streamRawPredict"}, {"x-litellm-api-key": f"Bearer {key}", "Authorization": f"Bearer {_CALLER_GOOGLE_TOKEN}"}, ) - spend: Final = {cell: _spend_for_key(key) for cell, key in keys.items()} + spend: Final = {cell: _spend_for_key(rig.proxy, key) for cell, key in keys.items()} standard: Final = _list_price(f"vertex_ai/{_CLAUDE_ROUTE_MODEL}") regional: Final = standard * _uplift(f"vertex_ai/{_CLAUDE_ROUTE_MODEL}") assert spend == pytest.approx( @@ -370,7 +369,7 @@ def test_v1beta_generate_content_on_a_vertex_deployment_bills_the_deployment_loc keys: Final = {cell: scenario.key() for cell in cells} for (deployment, action), key in keys.items(): _post_and_read(rig, f"/v1beta/models/{deployment}:{action}", body, {"x-goog-api-key": key}) - spend: Final = {cell: _spend_for_key(key) for cell, key in keys.items()} + spend: Final = {cell: _spend_for_key(rig.proxy, key) for cell, key in keys.items()} standard: Final = _list_price(f"vertex_ai/{_GEMINI}") regional: Final = standard * _uplift(f"vertex_ai/{_GEMINI}") assert spend == pytest.approx( @@ -404,8 +403,8 @@ def test_gemini_passthrough_bills_a_named_custom_priced_deployment_only_for_keys for key in (allowed[action], plain[action]): _post_and_read(rig, f"/gemini/v1beta/models/{_GEMINI}:{action}", body, {"x-goog-api-key": key}) spend: Final = { - **{("allowed", action): _spend_for_key(key) for action, key in allowed.items()}, - **{("plain", action): _spend_for_key(key) for action, key in plain.items()}, + **{("allowed", action): _spend_for_key(rig.proxy, key) for action, key in allowed.items()}, + **{("plain", action): _spend_for_key(rig.proxy, key) for action, key in plain.items()}, } custom_price: Final = _PROMPT_TOKENS * _CUSTOM_INPUT_RATE + _COMPLETION_TOKENS * _CUSTOM_OUTPUT_RATE list_price: Final = _list_price(f"gemini/{_GEMINI}") @@ -441,7 +440,10 @@ def test_openai_passthrough_stream_cut_off_after_its_usage_bills_what_the_finish assert response.status_code == 200 with pytest.raises(httpx.RemoteProtocolError): response.read() - spend: Final = {"finished": _spend_for_key(finished_key), "cut off": _spend_for_key(cut_key)} + spend: Final = { + "finished": _spend_for_key(rig.proxy, finished_key), + "cut off": _spend_for_key(rig.proxy, cut_key), + } list_price: Final = _list_price(_OPENAI_MODEL) assert spend == pytest.approx({"finished": list_price, "cut off": list_price}) @@ -484,7 +486,7 @@ def test_a_body_base_model_reprices_passthrough_spend_only_for_keys_that_allow_c }, _caller_headers(route, key), ) - spend: Final = {cell: _spend_for_key(key) for cell, key in keys.items()} + spend: Final = {cell: _spend_for_key(rig.proxy, key) for cell, key in keys.items()} forwarded: Final = rig.upstream.drain() marked_bodies: Final = { cell: tuple(request.body for request in forwarded if marker.encode() in request.body) @@ -521,7 +523,7 @@ def test_a_non_string_body_base_model_bills_the_list_price_instead_of_nothing(ri {"contents": [{"role": "user", "parts": [{"text": uuid.uuid4().hex}]}], "base_model": 123}, _caller_headers(route, key), ) - spend: Final = {cell: _spend_for_key(key) for cell, key in keys.items()} + spend: Final = {cell: _spend_for_key(rig.proxy, key) for cell, key in keys.items()} list_price: Final = { "vertex_ai": _list_price(f"vertex_ai/{_GEMINI}"), "gemini": _list_price(f"gemini/{_GEMINI}"), @@ -538,8 +540,8 @@ def test_gemini_passthrough_through_a_gemini_api_base_gateway_logs_and_bills_as_ keys: Final = {action: scenario.key() for action in _GEMINI_ACTIONS} for action, key in keys.items(): _post_and_read(gateway_rig, f"/gemini/v1beta/models/{_GEMINI}:{action}", body, {"x-goog-api-key": key}) - providers: Final = {action: _provider_for_key(key) for action, key in keys.items()} - spend: Final = {action: _spend_for_key(key) for action, key in keys.items()} + providers: Final = {action: _provider_for_key(gateway_rig.proxy, key) for action, key in keys.items()} + spend: Final = {action: _spend_for_key(gateway_rig.proxy, key) for action, key in keys.items()} forwarded: Final = tuple( path for path in (urlsplit(request.target).path for request in gateway_rig.upstream.drain()) @@ -570,8 +572,8 @@ def test_a_configured_passthrough_to_a_non_google_gateway_bills_as_vertex_whatev body, {"Authorization": f"Bearer {key}"}, ) - providers: Final = {cell: _provider_for_key(key) for cell, key in keys.items()} - spend: Final = {cell: _spend_for_key(key) for cell, key in keys.items()} + providers: Final = {cell: _provider_for_key(gateway_rig.proxy, key) for cell, key in keys.items()} + spend: Final = {cell: _spend_for_key(gateway_rig.proxy, key) for cell, key in keys.items()} short_path_spend: Final = {action: spend[("short gateway path", action)] for action in _GEMINI_ACTIONS} assert providers == {cell: "vertex_ai" for cell in cells} assert {action: spend[("vertex resource", action)] for action in _GEMINI_ACTIONS} == pytest.approx( diff --git a/tests/integration/spend/test_vertex_passthrough_regional_uplift_spend.py b/tests/integration/spend/test_vertex_passthrough_regional_uplift_spend.py index eaa8f52d4f4..0b2dd8910d6 100644 --- a/tests/integration/spend/test_vertex_passthrough_regional_uplift_spend.py +++ b/tests/integration/spend/test_vertex_passthrough_regional_uplift_spend.py @@ -10,7 +10,6 @@ import httpx import pytest import yaml from integration._support.client import Gateway, eventually, gateway_from_environment -from integration._support.database import read_rows from integration._support.forward_proxy import ForwardProxy, tunnelling_forward_proxy from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.tls import server_context, write_self_signed_cert @@ -150,12 +149,12 @@ def _spend(rig: _Rig, key: str, location: str, action: str, model: str = _MODEL) body: Final = response.read() assert response.status_code == 200, body call_id: Final = response.headers["x-litellm-call-id"] - return _logged_spend(call_id) + return _logged_spend(rig.proxy, call_id) -def _logged_spend(call_id: str) -> float: +def _logged_spend(gateway: Gateway, call_id: str) -> float: rows: Final = eventually( - lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)), + lambda: gateway.rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)), lambda found: len(found) == 1, seconds=70, ) @@ -193,7 +192,7 @@ def test_custom_priced_vertex_passthrough_bills_the_deployment_price_even_when_t call_id: Final = response.headers["x-litellm-call-id"] with pytest.raises(httpx.RemoteProtocolError): response.read() - cut_off: Final = _logged_spend(call_id) + cut_off: Final = _logged_spend(rig.proxy, call_id) custom_price: Final = _PROMPT_TOKENS * _CUSTOM_INPUT_RATE + _COMPLETION_TOKENS * _CUSTOM_OUTPUT_RATE assert {**finished, "cut off": cut_off} == pytest.approx( {_GENERATE: custom_price, _STREAM: custom_price, "cut off": custom_price}