scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(tests): read owned proxy rows from the owned proxy's database (#47333)

Browse source 
* 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
This commit is contained in:
Mateo Wang 2026-10-20 17:07:06 +00:00 • committed by GitHub
parent f3513bca9e
commit f925c067d2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 31 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
tests/integration/authorization/test_agent_least_privilege.py
View file

@ -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:

2
tests/integration/providers/test_hosted_vllm_stream_stall_wire.py
View file

@ -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)'

11
tests/integration/routing/test_responses_websocket_turn_admission.py
View file

@ -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,
)

38
tests/integration/spend/test_passthrough_stream_pricing_paths_spend.py
View file

@ -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(

9
tests/integration/spend/test_vertex_passthrough_regional_uplift_spend.py
View file

@ -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}