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}