scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-02 07:13:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 150

test(integration): fix four flaky integration tests (#49554)

Browse source 
Scope the Prisma migration lock test's advisory lock check to its own database, so a lock another xdist worker holds in its database no longer trips it

Keep in-process MCP peers' idle connections open for 30s instead of 1s, so the gateway's pooled HTTP client doesn't reuse a connection the peer is closing and hang the tool call

Delete database guardrails in the logging-only scope tests through DELETE /guardrails/{id} on the shared gateway instead of raw SQL, so a default-on guardrail stops applying to the next test on the same worker right away

Bound the wedged coordination Redis /customer/update check by the shortest Redis wait the proxy uses, like the burst check, so a slow runner can't fail it while a handler that waits on Redis still does

(cherry picked from commit 4f2c75ddf5)
This commit is contained in:
Yuneng Jiang 2026-10-31 21:13:30 -07:00
parent 9a929f871d
commit ba18de8a06
No known key found for this signature in database
6 changed files with 32 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
tests/integration/_support/asgi.py
View file

@ -15,7 +15,7 @@ from starlette.types import ASGIApp
@contextmanager
def asgi_server(
app: ASGIApp, *, before_stop: Callable[[], None] | None = None, port: int = 0, keepalive_seconds: int = 1
app: ASGIApp, *, before_stop: Callable[[], None] | None = None, port: int = 0, keepalive_seconds: int = 30
) -> Iterator[str]:
with socket.socket() as listener:
listener.bind(("127.0.0.1", port))

3
tests/integration/database/test_prisma_migration_lock.py
View file

@ -19,7 +19,8 @@ def _can_take_prismas_session_lock(connection: psycopg.Connection[tuple[object,
def _granted_advisory_keys(connection: psycopg.Connection[tuple[object, ...]]) -> frozenset[int]:
rows: Final = connection.execute(
"SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted"
"SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted "
"AND database = (SELECT oid FROM pg_database WHERE datname = current_database())"
).fetchall()
return frozenset(int(row[0]) for row in rows)

6
tests/integration/management/test_user_updates_wedged_coordination_redis.py
View file

@ -193,8 +193,10 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged(
timings["s6_customer_update"] = _timed_post(
candidate, "/customer/update", {"user_id": customer_id, "max_budget": 6.0}
)
assert timings["s6_customer_update"] < _HANDLER_BUDGET_SECONDS, (
f"/customer/update took {timings['s6_customer_update']:.3f}s with a wedged coordination Redis"
assert timings["s6_customer_update"] < _SHORTEST_REDIS_WAIT_SECONDS, (
f"/customer/update took {timings['s6_customer_update']:.3f}s, reaching the "
f"{_SHORTEST_REDIS_WAIT_SECONDS}s shortest wait the proxy puts on Redis, so it waited on the "
"wedged coordination Redis"
)
timings["s7_reset_spend"] = _timed_post(candidate, f"/key/{key}/reset_spend", {"reset_to": 0})
assert timings["s7_reset_spend"] < _HANDLER_BUDGET_SECONDS, (

13
tests/integration/observability/_logging_only_scope_support.py
View file

@ -22,7 +22,7 @@ import yaml
from anthropic import Anthropic, AsyncAnthropic
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from integration._support.client import Gateway, Scenario, eventually, object_value
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
from integration._support.database import read_rows, write_rows
from integration._support.upstream import delete_scenario, register_scenario
from integration._support.wire import Reply, Request, Wire
@ -792,6 +792,7 @@ def _insert_database_guardrail(
@contextmanager
def _database_guardrail(
gateway: Gateway,
identity: str,
policy_url: str,
scope: str | None,
@ -803,11 +804,15 @@ def _database_guardrail(
try:
yield
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
def _delete_database_guardrail(identity: str) -> None:
write_rows('DELETE FROM "LiteLLM_GuardrailsTable" WHERE guardrail_name=%s', (identity,))
def _delete_database_guardrail(gateway: Gateway, identity: str) -> None:
rows: Final = read_rows('SELECT guardrail_id FROM "LiteLLM_GuardrailsTable" WHERE guardrail_name=%s', (identity,))
responses: Final = tuple(
gateway.request("DELETE", f"/guardrails/{string_value(row['guardrail_id'])}") for row in rows
)
assert all(response.status_code == 200 for response in responses), [response.text for response in responses]
def _post_guardrail_body(

30
tests/integration/observability/test_logging_only_scope_config.py
View file

@ -112,7 +112,7 @@ def test_G7_database_invalid_scope_reads_preserve_pre_call_blocking(gateway: Gat
assert blocked_response.status_code == 400, blocked_response.text
assert _drain_upstream(gateway.upstream_url) == ()
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
finally:
delete_scenario(upstream_handle)
@ -670,7 +670,7 @@ def test_database_guardrail_load_keeps_pre_call_blocking(
for entry in entries
) == ((identity, "pre_call", "guardrail_intervened"),), entries
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
finally:
delete_scenario(upstream_handle)
@ -790,7 +790,7 @@ def test_G6_database_guardrail_polling_normalizes_invalid_scope_without_reinitia
) == ((identity, "pre_call", "guardrail_intervened"),), (call_id, entries)
assert _drain_upstream(gateway.upstream_url) == ()
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
finally:
delete_scenario(upstream_handle)
@ -943,7 +943,7 @@ def test_management_post_validates_logging_only_scope(
assert isinstance(after_list, list), after_list
assert identity not in {object_value(item)["guardrail_name"] for item in after_list}, after_list
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
def test_H5_management_put_rejection_preserves_database_and_runtime(gateway: Gateway, tmp_path: Path) -> None:
@ -1065,7 +1065,7 @@ def test_H5_management_put_rejection_preserves_database_and_runtime(gateway: Gat
for entry in entries
) == ((identity, "pre_call", "guardrail_intervened"),), (call_id, entries)
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
delete_scenario(upstream_handle)
@ -1144,7 +1144,7 @@ def test_H6_management_patch_clears_scope_when_switching_to_blocking_only_mode(
for entry in entries
) == ((identity, ("pre_call",), "guardrail_intervened"),), entries
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
delete_scenario(upstream_handle)
@ -1231,7 +1231,7 @@ def test_H7_management_patch_rejection_preserves_pre_call_guardrail(gateway: Gat
for entry in entries
) == ((identity, "pre_call", "guardrail_intervened"),), (call_id, entries)
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
delete_scenario(upstream_handle)
@ -1321,7 +1321,7 @@ def test_management_patch_default_on_heals_stored_scope(
for entry in entries
) == ((identity, "pre_call", "guardrail_intervened"),), entries
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
finally:
delete_scenario(upstream_handle)
@ -1418,7 +1418,7 @@ def test_H10_management_patch_null_scope_restores_both_directions(gateway: Gatew
(identity, "logging_only", "guardrail_intervened"),
), entries
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
delete_scenario(upstream_handle)
@ -1505,7 +1505,7 @@ def test_H11_management_patch_same_scope_is_idempotent_and_output_only(gateway:
(entry["guardrail_name"], entry["guardrail_mode"], entry["guardrail_status"]) for entry in entries
) == tuple((identity, "logging_only", "guardrail_intervened") for _ in expected_directions), entries
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
delete_scenario(upstream_handle)
@ -1538,7 +1538,7 @@ def test_H12_management_reads_expose_typed_logging_only_scope(gateway: Gateway,
assert object_value(info["litellm_params"])["logging_only_scope"] == "input", info
assert object_value(matches[0]["litellm_params"])["logging_only_scope"] == "input", matches
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
def test_H13_unauthenticated_management_and_chat_requests_do_not_scan(gateway: Gateway, tmp_path: Path) -> None:
@ -1617,7 +1617,7 @@ def test_K1_post_rejects_continue_flag_outside_logging_only(gateway: Gateway, tm
assert after_list == before_list, (before_list, after_list)
assert identity not in {object_value(item)["guardrail_name"] for item in after_list}, after_list
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
def test_K3_patch_flag_rejection_preserves_raw_stored_row(gateway: Gateway, tmp_path: Path) -> None:
@ -1663,7 +1663,7 @@ def test_K3_patch_flag_rejection_preserves_raw_stored_row(gateway: Gateway, tmp_
} == {key: after_params[key] for key in after_params if key != "api_key"}, (before_params, after_params)
assert str(after_params["api_key"]).startswith("litellm_enc::"), after_params
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
def test_K4_patch_dropping_logging_only_mode_clears_stored_flag(gateway: Gateway, tmp_path: Path) -> None:
@ -1701,7 +1701,7 @@ def test_K4_patch_dropping_logging_only_mode_clears_stored_flag(gateway: Gateway
assert params.get("logging_only_continue_on_input_failure") is None, persisted
assert params["mode"] == ["pre_call"], persisted
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
def test_K5_yaml_flag_on_blocking_guardrail_boots_and_keeps_blocking(gateway: Gateway, tmp_path: Path) -> None:
@ -1747,5 +1747,5 @@ def test_K5_yaml_flag_on_blocking_guardrail_boots_and_keeps_blocking(gateway: Ga
log_text: Final = owned.log.read_text()
assert "Ignoring logging_only_continue_on_input_failure" in log_text, log_text
finally:
_delete_database_guardrail(identity)
_delete_database_guardrail(gateway, identity)
delete_scenario(upstream_handle)

2
tests/integration/observability/test_logging_only_scope_runtime.py
View file

@ -1085,7 +1085,7 @@ def test_X2_scope_patch_toggles_during_concurrent_requests_keep_one_scan_per_res
with wire_server(policy) as edge:
config: Final = _empty_proxy_configuration(tmp_path, identity)
with (
_database_guardrail(identity, edge.url, "output", mode="logging_only"),
_database_guardrail(gateway, identity, edge.url, "output", mode="logging_only"),
owned_proxy(gateway, tmp_path, {}, config=config, workers=1) as candidate,
):
guardrails: Final = candidate.get("/v2/guardrails/list")["guardrails"]