test(integration): isolate shared proxy state behind six order-dependent flakes (#46172)

Each test below passed or failed on the same commit depending on test order
or port timing in the CircleCI integration workflow

- readiness probe only polls a port the launched proxy owns, so a stolen
  port no longer sends GET /health/integration-readiness to another test's
  wire server
- team update tests seed fresh member users, since a proxy admin relocation
  adds team members to the destination org and leaked org_b membership into
  the shared INTERNAL_USER
- realtime whisper tests wait until every proxy worker serves only the new
  gpt-realtime-whisper deployment, via a per-worker deployment id route
- configured fallbacks test runs on its own proxy and database, because
  /config/update persists router settings and cannot unset fallbacks, which
  left [] and changed the no-fallback error text other tests assert
- otel excluded services tests each get a private Redis so a sibling proxy
  cannot warm the budget registry cache the tenant span depends on
- Langfuse audit upstream uses unique response ids instead of len(targets),
  which could repeat under concurrency and drop SpendLogs rows
This commit is contained in:
yuneng-jiang 2026-10-13 09:55:11 -07:00 • committed by GitHub
parent 59ef56f636
commit ec58b889cc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 154 additions and 47 deletions

View file

@ -10,6 +10,7 @@ import sys
import time
import uuid
from collections.abc import Generator, Iterable, Iterator, Mapping
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager, suppress
from dataclasses import dataclass
from enum import Enum
@ -37,6 +38,7 @@ DB_PUSH: Final = ("--use_prisma_db_push",)
MIGRATE_DEPLOY: Final = ()
LEGACY_MIGRATE_DEPLOY: Final = ("--use_legacy_migration_resolver",)
CONFIGURATION: Final = TypeAdapter(dict[str, object])
_DEPLOYMENT_IDS: Final = TypeAdapter(dict[str, list[str]])
MCP_COVERAGE_CONFIG: Final = Path(__file__).resolve().parents[1] / "mcp_coverage.toml"
@ -208,6 +210,33 @@ class OwnedProxy:
identity: Final = response.headers.get("x-integration-worker-pid", "")
return int(identity) if response.status_code == 200 and identity.isdecimal() else None
def _deployment_ids_on_a_worker(self, model_name: str) -> tuple[int, frozenset[str]] | None:
try:
with httpx.Client(
base_url=str(self.gateway.client.base_url),
timeout=2,
trust_env=False,
limits=httpx.Limits(max_keepalive_connections=0),
) as client:
response: Final = client.get("/health/integration-deployment-ids", params={"model_name": model_name})
except httpx.TransportError:
return None
identity: Final = response.headers.get("x-integration-worker-pid", "")
if response.status_code != 200 or not identity.isdecimal():
return None
return int(identity), frozenset(_DEPLOYMENT_IDS.validate_python(response.json())["model_ids"])
def serves_only_on_every_worker(self, model_name: str, model_id: str) -> bool:
with ThreadPoolExecutor(max_workers=8) as pool:
answers: Final = tuple(pool.map(self._deployment_ids_on_a_worker, (model_name,) * 16))
samples: Final = tuple(answer for answer in answers if answer is not None)
workers: Final = self.worker_pids()
return (
bool(workers)
and workers <= frozenset(pid for pid, _ in samples)
and all(model_ids == {model_id} for _, model_ids in samples)
)
def _await_replacement(self, before: frozenset[int]) -> int:
budget: Final = readiness_seconds(1)
deadline: Final = time.monotonic() + budget
@ -494,6 +523,8 @@ def _prisma_running(launch: _Launch) -> bool:
def _answer(client: httpx.Client, launch: _Launch) -> tuple[bool, int | None]:
if not _owns_listener(launch):
return False, None
try:
response: Final = client.get("/health/integration-readiness", timeout=2)
except httpx.ConnectError:
@ -501,7 +532,7 @@ def _answer(client: httpx.Client, launch: _Launch) -> tuple[bool, int | None]:
except httpx.TransportError:
return True, None
identity: Final = response.headers.get("x-integration-worker-pid", "")
ready: Final = response.status_code == 200 and identity.isdecimal() and _owns_listener(launch)
ready: Final = response.status_code == 200 and identity.isdecimal()
return True, int(identity) if ready else None
@ -833,7 +864,7 @@ def refused_boot_log(
with httpx.Client(base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False) as client:
deadline: Final = time.monotonic() + readiness_seconds(1)
while launch.process.poll() is None:
assert not _is_ready(client), (
assert not (_owns_listener(launch) and _is_ready(client)), (
f"Proxy became ready instead of refusing to boot:\n{launch.log.read_text()}"
)
assert time.monotonic() < deadline, "Proxy neither exited nor became ready within the deadline"

View file

@ -26,7 +26,15 @@ async def worker_readiness() -> JSONResponse:
)
async def worker_deployment_ids(model_name: str) -> JSONResponse:
from litellm.proxy.proxy_server import llm_router
model_ids: Final = [] if llm_router is None else llm_router.get_model_ids(model_name=model_name)
return JSONResponse(content={"model_ids": model_ids}, headers={"x-integration-worker-pid": str(os.getpid())})
def configure_readiness() -> None:
from litellm.proxy.proxy_server import app
app.add_api_route("/health/integration-readiness", worker_readiness, include_in_schema=False)
app.add_api_route("/health/integration-deployment-ids", worker_deployment_ids, include_in_schema=False)

View file

@ -2,7 +2,7 @@ import pytest
from litellm.proxy._types import LitellmUserRoles
from tests.integration._support.actors import Actor
from tests.integration._support.actors import Actor, SeededKey
from tests.integration._support.actors import create_scratch_actor, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -38,19 +38,39 @@ _MATRIX = [
]
async def _seed_target(db, actors, shape: str, team_id: str) -> str:
_TEAM_MEMBER_ACTORS = (
Actor.TEAM_ADMIN,
Actor.INTERNAL_USER,
Actor.OWNER,
Actor.UNRELATED_SAME_ORG,
Actor.SERVICE_ACCOUNT,
Actor.CROSS_ORG_USER,
)
async def _callers(db, actors, prefix: str) -> dict[Actor, SeededKey]:
team_members = {
actor: await create_scratch_actor(
db, prefix, user_role=LitellmUserRoles.INTERNAL_USER.value, suffix=actor.value
)
for actor in _TEAM_MEMBER_ACTORS
}
return {**actors.keys, **team_members}
async def _seed_target(db, actors, callers, shape: str, team_id: str) -> str:
"""Raw-seed the resources target team; returns its organization_id."""
if shape == "alpha":
await create_scratch_team(
db,
team_id,
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
admin_user_ids=[callers[Actor.TEAM_ADMIN].user_id],
member_user_ids=[
actors.keys[Actor.INTERNAL_USER].user_id,
actors.keys[Actor.OWNER].user_id,
actors.keys[Actor.UNRELATED_SAME_ORG].user_id,
actors.keys[Actor.SERVICE_ACCOUNT].user_id,
callers[Actor.INTERNAL_USER].user_id,
callers[Actor.OWNER].user_id,
callers[Actor.UNRELATED_SAME_ORG].user_id,
callers[Actor.SERVICE_ACCOUNT].user_id,
],
)
return actors.org_a_id
@ -59,7 +79,7 @@ async def _seed_target(db, actors, shape: str, team_id: str) -> str:
db,
team_id,
organization_id=actors.org_b_id,
member_user_ids=[actors.keys[Actor.CROSS_ORG_USER].user_id],
member_user_ids=[callers[Actor.CROSS_ORG_USER].user_id],
)
return actors.org_b_id
pytest.fail(f"unknown shape={shape}") # pragma: no cover
@ -79,8 +99,9 @@ async def test_team_update_authz_matrix(
resources,
actors,
):
org_id = await _seed_target(db, actors, shape, resources.prefix)
caller = actors.keys[actor]
callers = await _callers(db, actors, resources.prefix)
org_id = await _seed_target(db, actors, callers, shape, resources.prefix)
caller = callers[actor]
resp = await gateway_client.post(
"/team/update",
@ -113,7 +134,8 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context(
(200), same as PROXY_ADMIN. A team admin of that same team reaches the
handler but is refused (403) until a proxy admin enables fields for team
admins, and the response says so."""
await _seed_target(db, actors, "alpha", resources.prefix)
callers = await _callers(db, actors, resources.prefix)
await _seed_target(db, actors, callers, "alpha", resources.prefix)
allowed_org_admin = await gateway_client.post(
"/team/update",
@ -131,7 +153,7 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context(
denied_team_admin = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {actors.keys[Actor.TEAM_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {callers[Actor.TEAM_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "team_alias": MARKER_ALIAS},
)
assert denied_team_admin.status_code == 403, denied_team_admin.text
@ -169,8 +191,9 @@ async def test_team_update_org_relocation_gate(
resources,
actors,
):
await _seed_target(db, actors, "alpha", resources.prefix)
caller = actors.keys[actor]
callers = await _callers(db, actors, resources.prefix)
await _seed_target(db, actors, callers, "alpha", resources.prefix)
caller = callers[actor]
resp = await gateway_client.post(
"/team/update",
@ -199,7 +222,7 @@ async def test_team_update_org_relocation_gate(
async def test_team_update_org_b_admin_relocation_rejection_detail(
gateway_client, db, resources, actors
):
await _seed_target(db, actors, "alpha", resources.prefix)
await _seed_target(db, actors, await _callers(db, actors, resources.prefix), "alpha", resources.prefix)
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {actors.keys[Actor.ORG_B_ADMIN].cleartext}"},

View file

@ -720,19 +720,17 @@ def _audit_upstream(provider_secret: str, marker: str, targets: list[str]):
assert request.headers["authorization"] == f"Bearer {provider_secret}"
targets.append(request.target)
body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(request.body)
index: Final = len(targets)
reply_id: Final = uuid.uuid4().hex
if request.target == "/v1/responses":
if body.get("stream"):
return Reply(
content_type="text/event-stream", chunks=_responses_stream_frames(f"resp-{marker}-{index}")
content_type="text/event-stream", chunks=_responses_stream_frames(f"resp-{marker}-{reply_id}")
)
return _responses_result(f"resp-{marker}-{index}")
return _responses_result(f"resp-{marker}-{reply_id}")
assert request.target == "/v1/chat/completions", request.target
if body.get("stream"):
return Reply(
content_type="text/event-stream", chunks=_chat_stream_frames(f"chatcmpl-{marker}-{index}")
)
return _completion(f"{marker}-{index}-answer")
return Reply(content_type="text/event-stream", chunks=_chat_stream_frames(f"chatcmpl-{marker}-{reply_id}"))
return _completion(f"{marker}-{reply_id}-answer")
return upstream

View file

@ -30,6 +30,7 @@ from integration._support.otlp_sink import (
spans_for_trace,
)
from integration._support.process import owned_proxy, owned_proxy_process
from integration._support.redis_process import owned_redis
from pydantic import JsonValue
DB_SYSTEM_KEYS: Final = frozenset({"db.system.name", "db.system"})
@ -42,6 +43,14 @@ def tenant_export_interval(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("OTEL_BSP_SCHEDULE_DELAY", str(TENANT_EXPORT_INTERVAL_MILLISECONDS))
@pytest.fixture(autouse=True)
def private_redis(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
with owned_redis(tmp_path / "redis") as cache:
monkeypatch.setenv("REDIS_HOST", cache.host)
monkeypatch.setenv("REDIS_PORT", str(cache.port))
yield
@pytest.fixture
def audit_sinks(tmp_path: Path) -> Iterator[SpanSinks]:
with owned_sinks(tmp_path / "sinks") as sinks:

View file

@ -362,7 +362,7 @@ def _muse_deployment(scenario: Scenario, scenario_id: str, api_base: str) -> str
return scenario.model(model=f"meta/{MUSE_MODEL}", api_key=scenario_id, api_base=api_base)
def _named_deployment(creator: Gateway, scenario: Scenario, name: str, model: str, scenario_id: str) -> str:
def _named_deployment_id(creator: Gateway, scenario: Scenario, name: str, model: str, scenario_id: str) -> str:
created: Final = creator.post(
"/model/new",
{
@ -371,10 +371,23 @@ def _named_deployment(creator: Gateway, scenario: Scenario, name: str, model: st
"model_info": {},
},
)
scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"]))
model_id: Final = string_value(object_value(created["model_info"])["id"])
scenario.cleanups.callback(scenario.delete_model, model_id)
return model_id
def _named_deployment(creator: Gateway, scenario: Scenario, name: str, model: str, scenario_id: str) -> str:
_named_deployment_id(creator, scenario, name, model, scenario_id)
return name
def _whisper_default_on_every_worker(owned: OwnedProxy, scenario: Scenario, scenario_id: str) -> None:
model_id: Final = _named_deployment_id(
owned.gateway, scenario, WHISPER_DEFAULT, f"openai/{WHISPER_DEFAULT}", scenario_id
)
eventually(lambda: owned.serves_only_on_every_worker(WHISPER_DEFAULT, model_id), lambda served: served, seconds=70)
def _close_code(closed: ConnectionClosed) -> int:
return 1006 if closed.rcvd is None else closed.rcvd.code
@ -868,9 +881,7 @@ def test_intent_without_model_routes_to_the_whisper_default_without_an_injected_
with guardrail_proxy.gateway.scenario() as scenario:
handle: Final = _scripted(scenario, _transcription_scenario(CLEAN_TRANSCRIPT))
key: Final = scenario.key()
_named_deployment(
guardrail_proxy.gateway, scenario, WHISPER_DEFAULT, f"openai/{WHISPER_DEFAULT}", handle.scenario_id
)
_whisper_default_on_every_worker(guardrail_proxy, scenario, handle.scenario_id)
session: Final = _transcribe(_ws_base(_owned_url(guardrail_proxy)), TRANSCRIPTION_QUERY, key)
observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)
forwarded: Final = _with_transcription_model(GA_TRANSCRIPTION_UPDATE, WHISPER_DEFAULT)
@ -1458,9 +1469,7 @@ def test_intent_without_model_blocked_transcript_reports_a_violation_on_the_whis
with guardrail_proxy.gateway.scenario() as scenario:
handle: Final = _scripted(scenario, _transcription_scenario(BLOCKED_TRANSCRIPT))
key: Final = scenario.key()
_named_deployment(
guardrail_proxy.gateway, scenario, WHISPER_DEFAULT, f"openai/{WHISPER_DEFAULT}", handle.scenario_id
)
_whisper_default_on_every_worker(guardrail_proxy, scenario, handle.scenario_id)
session: Final = _transcribe_blocked(_ws_base(_owned_url(guardrail_proxy)), TRANSCRIPTION_QUERY, key)
observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)
forwarded: Final = _with_transcription_model(GA_TRANSCRIPTION_UPDATE, WHISPER_DEFAULT)

View file

@ -42,7 +42,7 @@ from integration._support.client import (
object_value,
string_value,
)
from integration._support.database import read_rows
from integration._support.database import read_rows, scratch_database
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.upstream import delete_scenario, register_scenario
from integration._support.wire import Reply, Request, wire_server
@ -860,34 +860,34 @@ def test_hop_inside_one_encryption_boundary_keeps_the_reasoning(gateway: Gateway
assert [body.get("input") for body in bodies] == [history(mark), history(mark)], bodies
def configured_fallback(scenario: Scenario, source: str, target: str) -> None:
gateway: Final = scenario.gateway
original: Final = router_setting(gateway, "fallbacks")
restored: Final = original if isinstance(original, list) else []
scenario.cleanups.callback(lambda: gateway.post("/config/update", {"router_settings": {"fallbacks": restored}}))
gateway.post("/config/update", {"router_settings": {"fallbacks": [*restored, {source: [target]}]}})
def assert_cross_group_hop_stripped(gateway: Gateway, source_scenario: str, target_scenario: str, mark: str) -> None:
records: Final = observed(gateway)
assert [body.get("input") for body in bodies_for(records, source_scenario)] == [history(mark)], records
assert [body.get("input") for body in bodies_for(records, target_scenario)] == [stripped_history(mark)], records
def test_configured_fallbacks_entry_hop_strips_the_source_groups_reasoning(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
def test_configured_fallbacks_entry_hop_strips_the_source_groups_reasoning(gateway: Gateway, tmp_path: Path) -> None:
with (
scratch_database() as database,
owned_proxy_process(
gateway,
tmp_path,
{"DATABASE_URL": database},
remove_environment=("DIRECT_URL", "DATABASE_URL_READ_REPLICA"),
) as owned,
owned.gateway.scenario() as scenario,
):
source: Final = f"src-{uuid.uuid4().hex[:12]}"
target: Final = f"dst-{uuid.uuid4().hex[:12]}"
source_id: Final = deployment(scenario, source, scenario_api_base(scenario, source, failing()), order=None)
target_id: Final = deployment(scenario, target, scenario_api_base(scenario, target, healthy_json()), order=None)
configured_fallback(scenario, source, target)
owned.gateway.post("/config/update", {"router_settings": {"num_retries": 0, "fallbacks": [{source: [target]}]}})
mark: Final = marker()
answer: Final = responses_answer(
gateway, HTTPX, Target(source, frozenset({source_id, target_id})), history(mark)
owned.gateway, HTTPX, Target(source, frozenset({source_id, target_id})), history(mark)
)
assert_served_by(answer, target_id)
assert_cross_group_hop_stripped(gateway, source, target, mark)
assert source not in json.dumps(router_setting(gateway, "fallbacks")), router_setting(gateway, "fallbacks")
assert_cross_group_hop_stripped(owned.gateway, source, target, mark)
def test_client_fallbacks_entry_hop_strips_the_source_groups_reasoning(gateway: Gateway) -> None:

View file

@ -239,6 +239,35 @@ def test_only_live_readable_workers_can_satisfy_readiness(
process_module._wait_until_ready(launch)
def test_readiness_never_sends_a_request_to_a_port_the_launch_does_not_own(
process_module: ModuleType, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
monkeypatch.setattr(psutil, "process_iter", lambda: iter(()))
parent: Final = MagicMock()
parent.pid = 99
parent.poll.return_value = None
launch: Final = process_module._Launch(parent, 8123, tmp_path / "unused.log", 0.0)
original_client: Final = httpx.Client
def foreign_listener(request: httpx.Request) -> httpx.Response:
raise AssertionError(f"probed a listener another process owns: {request.url}")
def client(*, base_url: str, timeout: float, trust_env: bool, limits: httpx.Limits) -> httpx.Client:
return original_client(
base_url=base_url,
timeout=timeout,
trust_env=trust_env,
limits=limits,
transport=httpx.MockTransport(foreign_listener),
)
monkeypatch.setattr(httpx, "Client", client)
monkeypatch.setattr(time, "monotonic", lambda: 0.0)
monkeypatch.setenv("INTEGRATION_PROXY_READY_SECONDS", "-1")
with pytest.raises(AssertionError, match="readiness deadline"):
process_module._wait_until_ready(launch)
def test_boot_phases_mark_when_prisma_ran_the_port_bound_and_each_worker_answered(
process_module: ModuleType,
) -> None: