From ba18de8a06042042d6499e646332069cb9f168b7 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 31 Oct 2026 21:13:30 -0700 Subject: [PATCH] test(integration): fix four flaky integration tests (#49554) 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 4f2c75ddf539d1fa59737998ed022f7735d3bce5) --- tests/integration/_support/asgi.py | 2 +- .../database/test_prisma_migration_lock.py | 3 +- ..._user_updates_wedged_coordination_redis.py | 6 ++-- .../_logging_only_scope_support.py | 13 +++++--- .../test_logging_only_scope_config.py | 30 +++++++++---------- .../test_logging_only_scope_runtime.py | 2 +- 6 files changed, 32 insertions(+), 24 deletions(-) diff --git a/tests/integration/_support/asgi.py b/tests/integration/_support/asgi.py index e84dc9439d3..30eb9f21760 100644 --- a/tests/integration/_support/asgi.py +++ b/tests/integration/_support/asgi.py @@ -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)) diff --git a/tests/integration/database/test_prisma_migration_lock.py b/tests/integration/database/test_prisma_migration_lock.py index 7482e96d3d3..f4a2a022668 100644 --- a/tests/integration/database/test_prisma_migration_lock.py +++ b/tests/integration/database/test_prisma_migration_lock.py @@ -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) diff --git a/tests/integration/management/test_user_updates_wedged_coordination_redis.py b/tests/integration/management/test_user_updates_wedged_coordination_redis.py index b34f1a22e62..f2c10f9af22 100644 --- a/tests/integration/management/test_user_updates_wedged_coordination_redis.py +++ b/tests/integration/management/test_user_updates_wedged_coordination_redis.py @@ -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, ( diff --git a/tests/integration/observability/_logging_only_scope_support.py b/tests/integration/observability/_logging_only_scope_support.py index 24c5083d95a..222a833aef3 100644 --- a/tests/integration/observability/_logging_only_scope_support.py +++ b/tests/integration/observability/_logging_only_scope_support.py @@ -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( diff --git a/tests/integration/observability/test_logging_only_scope_config.py b/tests/integration/observability/test_logging_only_scope_config.py index 354e3b464fe..fe255613316 100644 --- a/tests/integration/observability/test_logging_only_scope_config.py +++ b/tests/integration/observability/test_logging_only_scope_config.py @@ -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) diff --git a/tests/integration/observability/test_logging_only_scope_runtime.py b/tests/integration/observability/test_logging_only_scope_runtime.py index b637859da52..bfd71cc6e1d 100644 --- a/tests/integration/observability/test_logging_only_scope_runtime.py +++ b/tests/integration/observability/test_logging_only_scope_runtime.py @@ -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"]