diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 6ec8dd7447a..a1e1b649202 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -40,9 +40,9 @@ CONFIGURATION: Final = TypeAdapter(dict[str, object]) MCP_COVERAGE_CONFIG: Final = Path(__file__).resolve().parents[1] / "mcp_coverage.toml" -def proxy_database_environment() -> Mapping[str, str]: - writer: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL", "") - reader: Final = os.environ.get("INTEGRATION_PROXY_READ_REPLICA_URL", "") +def proxy_database_environment(inherited: Mapping[str, str] = os.environ) -> Mapping[str, str]: + writer: Final = inherited.get("INTEGRATION_PROXY_DATABASE_URL", "") + reader: Final = inherited.get("INTEGRATION_PROXY_READ_REPLICA_URL", "") return MappingProxyType( { **({"DATABASE_URL": writer} if writer else {}), @@ -326,6 +326,7 @@ def owned_proxy( workers: int = 1, database_setup: tuple[str, ...] = DB_PUSH, coverage_file: Path | None = None, + inherited: Mapping[str, str] = os.environ, ) -> Iterator[Gateway]: with owned_proxy_process( gateway, @@ -336,6 +337,7 @@ def owned_proxy( workers=workers, database_setup=database_setup, coverage_file=coverage_file, + inherited=inherited, ) as owned: yield owned.gateway @@ -381,8 +383,13 @@ class _Launch: started: float +def _wall_clock(monotonic: float) -> float: + return time.time() - (time.monotonic() - monotonic) + + def _stop_launch(launch: _Launch) -> None: started: Final = time.monotonic() + stop_started_at: Final = time.time() exit_before: Final = launch.process.poll() exception_before: Final = sys.exc_info()[0] try: @@ -399,6 +406,9 @@ def _stop_launch(launch: _Launch) -> None: "exception_at_record": exception_at_record.__name__ if exception_at_record else None, "cleanup_seconds": time.monotonic() - started, "process_lifetime_seconds": time.monotonic() - launch.started, + "spawned_at": _wall_clock(launch.started), + "stop_started_at": stop_started_at, + "stopped_at": time.time(), }, indent=2, ) @@ -439,7 +449,51 @@ def _owns_listener(launch: _Launch) -> bool: return False -def _readiness_pids(launch: _Launch, workers: int) -> Iterator[int]: +@dataclass(frozen=True, slots=True) +class BootPhases: + prisma_started: float | None = None + prisma_finished: float | None = None + listening: float | None = None + first_ready: float | None = None + ready_pids: frozenset[int] = frozenset() + + +@dataclass(frozen=True, slots=True) +class _Probe: + elapsed: float + listening: bool + prisma: bool + ready_pid: int | None + + +def _runs_prisma(process: psutil.Process) -> bool: + try: + return any(Path(part).name == "prisma" for part in process.cmdline()) + except psutil.Error: + return False + + +def _prisma_running(launch: _Launch) -> bool: + try: + children: Final = psutil.Process(launch.process.pid).children() + except psutil.Error: + return False + return any(_runs_prisma(child) for child in children) + + +def _answer(client: httpx.Client, launch: _Launch) -> tuple[bool, int | None]: + try: + response: Final = client.get("/health/integration-readiness", timeout=2) + except httpx.ConnectError: + return False, 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) + return True, int(identity) if ready else None + + +def _probes(launch: _Launch, workers: int) -> Iterator[_Probe]: with httpx.Client( base_url=f"http://127.0.0.1:{launch.port}", timeout=15, @@ -448,19 +502,27 @@ def _readiness_pids(launch: _Launch, workers: int) -> Iterator[int]: ) as client: deadline: Final = time.monotonic() + readiness_seconds(workers) while launch.process.poll() is None: - try: - response: Final = client.get("/health/integration-readiness", timeout=2) - identity: Final = response.headers.get("x-integration-worker-pid", "") - if response.status_code == 200 and identity.isdecimal() and _owns_listener(launch): - yield int(identity) - except httpx.TransportError: - pass + elapsed: Final = time.monotonic() - launch.started + listening, ready_pid = _answer(client, launch) + yield _Probe(elapsed, listening, not listening and _prisma_running(launch), ready_pid) assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded" time.sleep(0.1) -def _ready_workers(pids: frozenset[int], pid: int) -> frozenset[int]: - return pids | frozenset((pid,)) +def _first(recorded: float | None, observed: bool, elapsed: float) -> float | None: + return recorded if recorded is not None or not observed else elapsed + + +def _observe(phases: BootPhases, probe: _Probe) -> BootPhases: + prisma_started: Final = _first(phases.prisma_started, probe.prisma, probe.elapsed) + prisma_finished: Final = _first(phases.prisma_finished, prisma_started is not None, probe.elapsed) + return BootPhases( + prisma_started=prisma_started, + prisma_finished=None if probe.prisma else prisma_finished, + listening=_first(phases.listening, probe.listening, probe.elapsed), + first_ready=_first(phases.first_ready, probe.ready_pid is not None, probe.elapsed), + ready_pids=phases.ready_pids | frozenset(() if probe.ready_pid is None else (probe.ready_pid,)), + ) def _ready_worker_pid(member: psutil.Process) -> int | None: @@ -475,13 +537,24 @@ def _ready_worker_pid(member: psutil.Process) -> int | None: return None -def _wait_until_ready(launch: _Launch, workers: int = 1) -> None: - for ready in accumulate(_readiness_pids(launch, workers), _ready_workers, initial=frozenset[int]()): - members: Final = frozenset( - pid for member in group_members(launch.process.pid) if (pid := _ready_worker_pid(member)) is not None - ) - if len(ready & members) >= max(1, workers): - return +def _serves_all(launch: _Launch, phases: BootPhases, workers: int) -> bool: + if len(phases.ready_pids) < max(1, workers): + return False + members: Final = frozenset( + pid for member in group_members(launch.process.pid) if (pid := _ready_worker_pid(member)) is not None + ) + return len(phases.ready_pids & members) >= max(1, workers) + + +def _wait_until_ready(launch: _Launch, workers: int = 1) -> BootPhases | None: + return next( + ( + phases + for phases in accumulate(_probes(launch, workers), _observe, initial=BootPhases()) + if _serves_all(launch, phases, workers) + ), + None, + ) def _launch_until_bound( @@ -489,7 +562,7 @@ def _launch_until_bound( ) -> _Launch: launch: Final = _launch(command, root, environment, output) try: - _wait_until_ready(launch, workers) + phases: Final = _wait_until_ready(launch, workers) exit_code: Final = launch.process.poll() assert exit_code is None or (attempts > 1 and _lost_port_race(exit_code, launch.log)), ( "Owned proxy exited before readiness" @@ -498,12 +571,21 @@ def _launch_until_bound( _stop_launch(launch) raise if exit_code is None: + assert phases is not None launch.log.with_suffix(".startup.json").write_text( json.dumps( { "root_pid": launch.process.pid, "readiness_seconds": time.monotonic() - launch.started, "logical_cpus": os.cpu_count(), + "workers": workers, + "spawned_at": _wall_clock(launch.started), + "phases": { + "prisma_started": phases.prisma_started, + "prisma_finished": phases.prisma_finished, + "listening": phases.listening, + "first_ready": phases.first_ready, + }, } ) + "\n" @@ -517,14 +599,14 @@ def _proxy_root() -> Path: return Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) -def worker_configuration(config: Path, directory: Path) -> Path: - if os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1": +def worker_configuration(config: Path, directory: Path, inherited: Mapping[str, str] = os.environ) -> Path: + if inherited.get("INTEGRATION_WORKER_ISOLATION") != "1": return config def endpoint(match: re.Match[str]) -> str: if match.group(1) == "8190": - return os.environ["INTEGRATION_UPSTREAM_URL"] - return f"http://127.0.0.1:{os.environ['INTEGRATION_PROVIDER_PORT']}" + return inherited["INTEGRATION_UPSTREAM_URL"] + return f"http://127.0.0.1:{inherited['INTEGRATION_PROVIDER_PORT']}" contents: Final = re.sub(r"http://127[.]0[.]0[.]1:(8190|8191)(?![0-9])", endpoint, config.read_text()) directory.mkdir(parents=True, exist_ok=True) @@ -547,9 +629,12 @@ def worker_configuration(config: Path, directory: Path) -> Path: def _proxy_environment( - gateway: Gateway, overrides: Mapping[str, str], remove_environment: tuple[str, ...] + gateway: Gateway, + overrides: Mapping[str, str], + remove_environment: tuple[str, ...], + source: Mapping[str, str] = os.environ, ) -> Mapping[str, str]: - inherited: Final = {**os.environ, **proxy_database_environment()} + inherited: Final = {**source, **proxy_database_environment(source)} writer: Final = overrides.get("DATABASE_URL") reader: Final = inherited.get("DATABASE_URL_READ_REPLICA") unpaired_reader: Final = ( @@ -567,7 +652,7 @@ def _proxy_environment( { **{name: value for name, value in inherited.items() if name not in removed}, "LITELLM_MASTER_KEY": gateway.key, - "LITELLM_SALT_KEY": os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt"), + "LITELLM_SALT_KEY": source.get("LITELLM_SALT_KEY", "sk-integration-salt"), "STORE_MODEL_IN_DB": "True", "PYTHON_DOTENV_DISABLED": "1", **overrides, @@ -628,13 +713,17 @@ def owned_proxy_process( extra_arguments: tuple[str, ...] = (), coverage_file: Path | None = None, healthcheck: WorkerHealthcheck = WorkerHealthcheck.STRICT, + inherited: Mapping[str, str] = os.environ, ) -> Iterator[OwnedProxy]: assert "--timeout_worker_healthcheck" not in extra_arguments, "Choose the worker healthcheck with `healthcheck`" root: Final = _proxy_root() environment: Final = _proxy_environment( - gateway, {**overrides, **({"COVERAGE_FILE": str(coverage_file)} if coverage_file else {})}, remove_environment + gateway, + {**overrides, **({"COVERAGE_FILE": str(coverage_file)} if coverage_file else {})}, + remove_environment, + inherited, ) - output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory))) + output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory))) output.mkdir(parents=True, exist_ok=True) command: Final = ( sys.executable, @@ -642,7 +731,7 @@ def owned_proxy_process( *(("coverage", "run", f"--rcfile={MCP_COVERAGE_CONFIG}", "-m") if coverage_file else ()), "integration._support.proxy", "--config", - str(worker_configuration(config or root / "tests/integration/proxy_config.yaml", directory)), + str(worker_configuration(config or root / "tests/integration/proxy_config.yaml", directory, inherited)), "--host", "127.0.0.1", "--num_workers", @@ -786,13 +875,21 @@ class UpstreamSlot: It runs from the harness's own checkout, never ``INTEGRATION_PROXY_ROOT``: the double belongs to the tests, the proxy root only names the litellm under test.""" - __slots__ = ("certificate", "directory", "port", "process", "root") + __slots__ = ("certificate", "directory", "inherited", "port", "process", "root") - def __init__(self, directory: Path, port: int, root: Path, certificate: UpstreamCertificate | None = None) -> None: + def __init__( + self, + directory: Path, + port: int, + root: Path, + certificate: UpstreamCertificate | None = None, + inherited: Mapping[str, str] = os.environ, + ) -> None: self.directory = directory self.port = port self.root = root self.certificate = certificate + self.inherited = inherited self.process: subprocess.Popen[bytes] | None = None @property @@ -807,7 +904,7 @@ class UpstreamSlot: def start(self) -> None: assert self.process is None, "Owned upstream is already running" - output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR") or self.directory) + output: Final = Path(self.inherited.get("INTEGRATION_RESULTS_DIR") or self.directory) log_path: Final = output / f"owned-upstream-{self.port}-{uuid.uuid4().hex}.log" with log_path.open("w") as log: process: Final = subprocess.Popen( @@ -820,7 +917,7 @@ class UpstreamSlot: *self._tls_arguments(), ], cwd=self.root, - env=dict(os.environ), + env=dict(self.inherited), stdout=log, stderr=subprocess.STDOUT, start_new_session=True, @@ -856,9 +953,11 @@ def _harness_root() -> Path: @contextmanager -def owned_upstream(directory: Path, *, tls: bool = False, port: int | None = None) -> Generator[UpstreamSlot]: +def owned_upstream( + directory: Path, *, tls: bool = False, port: int | None = None, inherited: Mapping[str, str] = os.environ +) -> Generator[UpstreamSlot]: certificate: Final = self_signed_certificate(directory) if tls else None - slot: Final = UpstreamSlot(directory, port or _free_port(), _harness_root(), certificate) + slot: Final = UpstreamSlot(directory, port or _free_port(), _harness_root(), certificate, inherited) try: slot.start() yield slot diff --git a/tests/integration/_support/redis_process.py b/tests/integration/_support/redis_process.py index cfdbf041eba..f5e7f509eb5 100644 --- a/tests/integration/_support/redis_process.py +++ b/tests/integration/_support/redis_process.py @@ -5,7 +5,7 @@ import socket import subprocess import time import uuid -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path @@ -172,7 +172,11 @@ def _started( @contextmanager -def owned_redis(directory: Path, choose_port: Callable[[], int] = free_loopback_port) -> Iterator[OwnedRedis]: +def owned_redis( + directory: Path, + choose_port: Callable[[], int] = free_loopback_port, + inherited: Mapping[str, str] = os.environ, +) -> Iterator[OwnedRedis]: binary: Final = shutil.which("redis-server") if binary: host: Final = "127.0.0.1" @@ -190,7 +194,7 @@ def owned_redis(directory: Path, choose_port: Callable[[], int] = free_loopback_ ).strip() assert host, "CircleCI owned Redis container has no address" prefix: Final = ("docker", "exec", "redis-cache", "redis-server") - output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory))) + output: Final = Path(inherited.get("INTEGRATION_RESULTS_DIR", str(directory))) output.mkdir(parents=True, exist_ok=True) directory.mkdir(parents=True, exist_ok=True) with (output / f"owned-redis-{uuid.uuid4().hex}.log").open("w") as log: diff --git a/tests/integration/_support/workers.py b/tests/integration/_support/workers.py index 843f5dec4ac..33e086576c5 100644 --- a/tests/integration/_support/workers.py +++ b/tests/integration/_support/workers.py @@ -4,8 +4,9 @@ import os import socket import time import uuid -from collections.abc import Iterator, Mapping -from contextlib import contextmanager, nullcontext +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack, contextmanager, nullcontext from dataclasses import dataclass from pathlib import Path from types import MappingProxyType @@ -151,11 +152,13 @@ def cost_map_environment(upstream: str) -> Mapping[str, str]: @contextmanager -def worker_services(directory: Path, *, replica: bool = False) -> Iterator[Mapping[str, str]]: +def worker_services( + directory: Path, *, replica: bool = False, inherited: Mapping[str, str] = os.environ +) -> Iterator[Mapping[str, str]]: directory.mkdir(parents=True, exist_ok=True) with ( scratch_database(WORKER_DATABASE_PREFIX) as database, - owned_redis(directory) as cache, + owned_redis(directory, inherited=inherited) as cache, ): yield MappingProxyType( { @@ -170,15 +173,17 @@ def worker_services(directory: Path, *, replica: bool = False) -> Iterator[Mappi @contextmanager -def worker_gateway(directory: Path, profile: GatewayProfile) -> Iterator[Mapping[str, str]]: +def worker_gateway( + directory: Path, profile: GatewayProfile, inherited: Mapping[str, str] +) -> Iterator[Mapping[str, str]]: from tests.integration._support.process import owned_proxy, owned_upstream - upstream: Final = os.environ["INTEGRATION_UPSTREAM_URL"] + upstream: Final = inherited["INTEGRATION_UPSTREAM_URL"] port: Final = urlsplit(upstream).port assert port is not None key: Final = f"sk-{uuid.uuid4().hex}" removed: Final = ("LITELLM_LOCAL_MODEL_COST_MAP",) if profile.cost_map else () - with worker_services(directory, replica=profile.replica) as services: + with worker_services(directory, replica=profile.replica, inherited=inherited) as services: overrides: Final = { **services, "DATABASE_URL": services["INTEGRATION_PROXY_DATABASE_URL"], @@ -189,7 +194,7 @@ def worker_gateway(directory: Path, profile: GatewayProfile) -> Iterator[Mapping **(cost_map_environment(upstream) if profile.cost_map else {}), } with ( - owned_upstream(directory, port=port), + owned_upstream(directory, port=port, inherited=inherited), httpx.Client(base_url="http://127.0.0.1:1", trust_env=False) as client, owned_proxy( Gateway(client, key, upstream), @@ -197,8 +202,16 @@ def worker_gateway(directory: Path, profile: GatewayProfile) -> Iterator[Mapping overrides, remove_environment=removed, coverage_file=profile.coverage_file, + inherited=inherited, ) as primary, - owned_proxy(primary, directory, overrides, remove_environment=removed, coverage_file=profile.coverage_file) + owned_proxy( + primary, + directory, + overrides, + remove_environment=removed, + coverage_file=profile.coverage_file, + inherited=inherited, + ) if profile.peer else nullcontext(None) as secondary, ): @@ -215,3 +228,36 @@ def worker_gateway(directory: Path, profile: GatewayProfile) -> Iterator[Mapping "INTEGRATION_UPSTREAM_URL": upstream, } ) + + +@dataclass(frozen=True, slots=True) +class WorkerStack: + resources: ExitStack + environment: Mapping[str, str] + + +def worker_stack(results: Path, profile: GatewayProfile) -> WorkerStack: + endpoints: Final = worker_endpoints() + inherited: Final = MappingProxyType({**os.environ, **endpoints, "INTEGRATION_RESULTS_DIR": str(results)}) + resources: Final = ExitStack() + try: + services: Final = resources.enter_context(worker_gateway(results / "runtime", profile, inherited)) + except BaseException: + resources.close() + raise + return WorkerStack(resources, MappingProxyType({**endpoints, **services, "INTEGRATION_RESULTS_DIR": str(results)})) + + +def start_worker_stacks( + directories: tuple[Path, ...], build: Callable[[Path], WorkerStack], resources: ExitStack, concurrency: int +) -> tuple[Mapping[str, str], ...]: + with ThreadPoolExecutor(max_workers=max(1, concurrency), thread_name_prefix="integration-worker-stack") as pool: + futures: Final = tuple(pool.submit(build, directory) for directory in directories) + for future in futures: + if future.exception() is None: + resources.enter_context(future.result().resources) + failure: Final = next((error for future in futures if (error := future.exception()) is not None), None) + if failure is not None: + resources.close() + raise failure + return tuple(future.result().environment for future in futures) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 8bbf4e71842..cb23713e6b2 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -7,6 +7,7 @@ import re from collections.abc import AsyncIterator, Generator, Iterator, Mapping, Sequence from contextlib import ExitStack from copy import copy +from functools import partial from importlib.metadata import version from pathlib import Path from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable @@ -24,7 +25,7 @@ from tests.integration._support.manifest import OWNED_DIRECTORIES from tests.integration._support.ordering import collected_case, fixture_groups, scheduling_groups from tests.integration._support.provider import SharedProvider, shared_provider from tests.integration._support.routing import RoutingPlugin -from tests.integration._support.workers import gateway_profile, worker_endpoints, worker_gateway +from tests.integration._support.workers import gateway_profile, start_worker_stacks, worker_stack from tests.integration.run import GITHUB_FILES if TYPE_CHECKING: @@ -84,21 +85,14 @@ def pytest_xdist_setupnodes(config: pytest.Config, specs: Sequence[XSpec]) -> No resources: Final = ExitStack() config.add_cleanup(resources.close) destination: Final = Path(os.environ["INTEGRATION_RESULTS_DIR"]) - profile: Final = gateway_profile(os.environ) - for index, spec in enumerate(specs): - results: Final = destination / f"gw{index}" - endpoints: Final = worker_endpoints() - with pytest.MonkeyPatch.context() as environment: - for name, value in endpoints.items(): - environment.setenv(name, value) - environment.setenv("INTEGRATION_RESULTS_DIR", str(results)) - services: Final = resources.enter_context(worker_gateway(results / "runtime", profile)) - spec.env = { - **STRING_ENVIRONMENT.validate_python(spec.env), - **endpoints, - **services, - "INTEGRATION_RESULTS_DIR": str(results), - } + stacks: Final = start_worker_stacks( + tuple(destination / f"gw{index}" for index in range(len(specs))), + partial(worker_stack, profile=gateway_profile(os.environ)), + resources, + os.cpu_count() or 1, + ) + for spec, stack in zip(specs, stacks, strict=True): + spec.env = {**STRING_ENVIRONMENT.validate_python(spec.env), **stack} class IntegrationReportPlugin: diff --git a/tests/integration/observability/test_otel_excluded_services.py b/tests/integration/observability/test_otel_excluded_services.py index 1de071e7103..d62537e2317 100644 --- a/tests/integration/observability/test_otel_excluded_services.py +++ b/tests/integration/observability/test_otel_excluded_services.py @@ -3,6 +3,8 @@ from __future__ import annotations import time import uuid from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass, replace +from itertools import accumulate from pathlib import Path from types import MappingProxyType from typing import Final @@ -29,6 +31,13 @@ from integration._support.process import owned_proxy, owned_proxy_process from pydantic import JsonValue DB_SYSTEM_KEYS: Final = frozenset({"db.system.name", "db.system"}) +TENANT_EXPORT_INTERVAL_MILLISECONDS: Final = 1000 +TAIL_QUIET_SECONDS: Final = 5 * TENANT_EXPORT_INTERVAL_MILLISECONDS / 1000 + + +@pytest.fixture(autouse=True) +def tenant_export_interval(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OTEL_BSP_SCHEDULE_DELAY", str(TENANT_EXPORT_INTERVAL_MILLISECONDS)) @pytest.fixture @@ -125,21 +134,33 @@ def _trace_id(sink_url: str, response: httpx.Response, seconds: float = 40) -> s return found -def _trace_spans(sink_url: str, trace_id: str, seconds: float = 30) -> tuple[Span, ...]: - """The trace's spans once the post-call tail has landed. +@dataclass(frozen=True, slots=True) +class _TraceSample: + at: float + changed_at: float + spans: tuple[Span, ...] - The spend-writer and other post-response spans flush after the request - answers, so absence assertions poll for the whole window instead of - settling at the first glimpse of the root span. - """ - deadline: Final = time.monotonic() + seconds - group: tuple[Span, ...] = () # rebind-ok: drains samples until the post-call tail lands - while time.monotonic() < deadline: - _, spans = recorded_spans(sink_url) - group = spans_for_trace(spans, trace_id) + +def _trace_samples(sink_url: str, trace_id: str) -> Iterator[_TraceSample]: + while True: + now: Final = time.monotonic() + yield _TraceSample(now, now, spans_for_trace(recorded_spans(sink_url)[1], trace_id)) time.sleep(0.5) - assert group, f"trace {trace_id} never reached {sink_url}" - return group + + +def _keep_unchanged_since(previous: _TraceSample, current: _TraceSample) -> _TraceSample: + return current if len(current.spans) != len(previous.spans) else replace(current, changed_at=previous.changed_at) + + +def _trace_spans(sink_url: str, trace_id: str, seconds: float = 30) -> tuple[Span, ...]: + deadline: Final = time.monotonic() + seconds + settled: Final = next( + sample + for sample in accumulate(_trace_samples(sink_url, trace_id), _keep_unchanged_since) + if (sample.spans and sample.at - sample.changed_at >= TAIL_QUIET_SECONDS) or sample.at >= deadline + ) + assert settled.spans, f"trace {trace_id} never reached {sink_url}" + return settled.spans def _trace_spans_when( @@ -231,7 +252,7 @@ def test_excluded_services_drops_db_spans_at_tenant_only( config: Final = _config_with( tmp_path, otel_audit_config, otel={"excluded_services": ["redis", "postgres"]}, extra=_guardrail_block ) - with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config) as candidate: ten_start, _ = recorded_spans(audit_sinks.tenant) op_start, _ = recorded_spans(audit_sinks.operator) traffic: Final = _drive(candidate, langfuse_vars) @@ -266,7 +287,7 @@ def test_without_excluded_services_the_tenant_still_gets_redis_and_postgres_span tmp_path: Path, ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, extra=_guardrail_block) - with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config) as candidate: tenant_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic) @@ -294,7 +315,7 @@ def test_a_non_mapping_otel_block_still_publishes_the_tenant_fan_out( config: Final = _config_with(tmp_path, otel_audit_config, extra=with_callback_settings) overrides: Final = {"LITELLM_OTEL_V2": "1", **_operator_langfuse(audit_sinks)} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: traffic: Final = _drive(candidate, langfuse_vars) tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic) _await_db_span(audit_sinks.tenant, tenant_trace, "redis") @@ -328,7 +349,7 @@ def test_a_bare_excluded_services_env_var_is_ignored( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config) overrides: Final = {"LITELLM_OTEL_V2": "1", name: "redis,postgres"} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: tenant_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic) @@ -360,7 +381,7 @@ def test_the_documented_env_var_wins_over_a_bare_excluded_services( "LITELLM_OTEL_EXCLUDED_SERVICES": "redis", "EXCLUDED_SERVICES": "postgres", } - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: tenant_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) tenant_trace: Final = _trace_id(audit_sinks.tenant, traffic) @@ -397,7 +418,6 @@ def test_case_sensitive_otel_settings_read_only_the_exact_env_name( overrides, config=config, remove_environment=("LITELLM_OTEL_EXCLUDED_SERVICES", "litellm_otel_excluded_services"), - workers=2, ) as candidate: tenant_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) @@ -431,7 +451,7 @@ def test_env_ignore_empty_keeps_the_default_service_name( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"_env_ignore_empty": True}) overrides: Final = {"LITELLM_OTEL_V2": "1", "OTEL_SERVICE_NAME": ""} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: traffic: Final = _drive(candidate, langfuse_vars) operator_trace: Final = _trace_id(audit_sinks.operator, traffic) operator_spans: Final = _trace_spans_when( @@ -456,7 +476,7 @@ def test_env_parse_none_str_reads_a_null_traces_endpoint_as_unset( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"_env_parse_none_str": "null"}) overrides: Final = {"LITELLM_OTEL_V2": "1", "OTEL_TRACES_ENDPOINT": "null"} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: traffic: Final = _drive(candidate, langfuse_vars) operator_trace: Final = _trace_id(audit_sinks.operator, traffic) operator_spans: Final = _trace_spans_when( @@ -477,7 +497,7 @@ def test_env_excluded_services_drops_only_redis( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config) with owned_proxy( - gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config, workers=2 + gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config ) as candidate: start, _ = recorded_spans(audit_sinks.tenant) _drive(candidate, langfuse_vars) @@ -503,7 +523,7 @@ def test_config_excluded_services_wins_over_env( tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}, extra=with_langfuse_otel ) with owned_proxy( - gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config, workers=2 + gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "redis"}, config=config ) as candidate: _assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars) @@ -523,7 +543,7 @@ def test_excluded_services_applies_with_preset_ordered_first( tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}, extra=preset_first ) overrides: Final = {"LITELLM_OTEL_V2": "1", **_operator_langfuse(audit_sinks)} - with owned_proxy(gateway, tmp_path, overrides, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate: _assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars) @@ -536,7 +556,7 @@ def test_bogus_excluded_service_logs_error_and_drops_at_proxy_start( tmp_path: Path, ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["auth", "postgres"]}) - with owned_proxy_process(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config) as owned: assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:] _assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars) @@ -551,7 +571,7 @@ def test_valid_config_excluded_services_tolerates_bogus_env( ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}) with owned_proxy( - gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth"}, config=config, workers=2 + gateway, tmp_path, {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth"}, config=config ) as candidate: _assert_tenant_keeps_redis_without_postgres(candidate, audit_sinks, langfuse_vars) @@ -569,7 +589,7 @@ def test_bogus_excluded_services_env_logs_and_drops_with_preset_alongside_otel( config: Final = _config_with(tmp_path, otel_audit_config, extra=with_langfuse_otel) overrides: Final = {"LITELLM_OTEL_V2": "1", "LITELLM_OTEL_EXCLUDED_SERVICES": "auth,postgres"} - with owned_proxy_process(gateway, tmp_path, overrides, config=config, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned: assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:] _assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars) @@ -591,7 +611,7 @@ def test_bogus_excluded_services_env_logs_and_drops_without_otel_callback( "LITELLM_OTEL_EXCLUDED_SERVICES": "auth,postgres", **_operator_langfuse(audit_sinks), } - with owned_proxy_process(gateway, tmp_path, overrides, config=config, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned: assert "'auth' is not a datastore service; ignored" in owned.log.read_text(), owned.log.read_text()[-3000:] _assert_tenant_keeps_redis_without_postgres(owned.gateway, audit_sinks, langfuse_vars) @@ -604,7 +624,7 @@ def test_postgres_exclusion_covers_spend_flush( tmp_path: Path, ) -> None: config: Final = _config_with(tmp_path, otel_audit_config, otel={"excluded_services": ["postgres"]}) - with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config, workers=2) as candidate: + with owned_proxy(gateway, tmp_path, {"LITELLM_OTEL_V2": "1"}, config=config) as candidate: op_start, _ = recorded_spans(audit_sinks.operator) ten_start, _ = recorded_spans(audit_sinks.tenant) traffic: Final = _drive(candidate, langfuse_vars) diff --git a/tests/integration/observability/test_passthrough_upstream_error_visibility.py b/tests/integration/observability/test_passthrough_upstream_error_visibility.py index bc6747844ba..0d86440045b 100644 --- a/tests/integration/observability/test_passthrough_upstream_error_visibility.py +++ b/tests/integration/observability/test_passthrough_upstream_error_visibility.py @@ -126,7 +126,7 @@ async def test_gemini_passthrough_async_client_404_body_reaches_proxy_log_and_sp path: Final = tmp_path / "gemini-async.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway async with httpx.AsyncClient( base_url=str(candidate.client.base_url), timeout=15, trust_env=False @@ -156,7 +156,7 @@ def test_gemini_passthrough_streaming_500_relays_full_body_and_logs_bounded_prev path: Final = tmp_path / "gemini-stream-500.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -188,7 +188,7 @@ def test_gemini_passthrough_success_logs_nothing_and_spend_row_is_success(gatewa path: Final = tmp_path / "gemini-200.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -208,7 +208,7 @@ def test_gemini_passthrough_streaming_200_relays_every_chunk(gateway: Gateway, t path: Final = tmp_path / "gemini-stream-200.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -239,7 +239,7 @@ def test_gemini_streaming_passthrough_request_tags_reach_spend_log(gateway: Gate config_path: Final = tmp_path / "gemini-stream-tags.yaml" with wire_server(respond) as wire: _gemini_config(config_path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=config_path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=config_path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -293,7 +293,7 @@ def test_config_pass_through_route_logs_body_and_strips_query(gateway: Gateway, } ] path.write_text(yaml.safe_dump(config)) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request("POST", "/audit-pt", _GENERATE_CONTENT) assert response.status_code == 403, response.text @@ -324,7 +324,7 @@ def test_openai_passthrough_sdk_error_body_reaches_proxy_log_and_spend_row(gatew path: Final = tmp_path / "openai-404.yaml" with wire_server(respond) as wire: _openai_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with OpenAI( api_key=candidate.key, @@ -350,7 +350,7 @@ async def test_openai_passthrough_async_sdk_error_body_reaches_proxy_log_and_spe path: Final = tmp_path / "openai-async-404.yaml" with wire_server(respond) as wire: _openai_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway async with AsyncOpenAI( api_key=candidate.key, @@ -377,7 +377,7 @@ def test_gemini_passthrough_control_characters_cannot_forge_log_lines(gateway: G path: Final = tmp_path / "gemini-forged.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -401,7 +401,7 @@ def test_gemini_passthrough_empty_error_body_still_logged_and_proxy_serves(gatew path: Final = tmp_path / "gemini-empty.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -434,7 +434,7 @@ def test_gemini_passthrough_gzip_error_body_decoded_for_log_and_client(gateway: path: Final = tmp_path / "gemini-gzip.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -462,7 +462,7 @@ def test_gemini_passthrough_streaming_gzip_error_body_decoded_for_log_and_client path: Final = tmp_path / "gemini-stream-gzip.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -490,7 +490,7 @@ def test_gemini_passthrough_error_body_redacted_when_message_logging_off(gateway config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"} config["litellm_settings"]["turn_off_message_logging"] = True path.write_text(yaml.safe_dump(config)) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -516,7 +516,7 @@ def test_gemini_passthrough_exact_4096_byte_body_logged_without_marker(gateway: path: Final = tmp_path / "gemini-exact.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -537,7 +537,7 @@ def test_gemini_passthrough_4097_byte_body_truncated_with_marker(gateway: Gatewa path: Final = tmp_path / "gemini-over.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway response: Final = candidate.request( "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} @@ -557,7 +557,7 @@ def test_gemini_passthrough_one_byte_stream_chunks_reassembled_and_logged(gatewa path: Final = tmp_path / "gemini-one-byte.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.client.stream( "POST", @@ -580,7 +580,7 @@ def test_gemini_passthrough_repeated_errors_each_get_row_and_log_line(gateway: G path: Final = tmp_path / "gemini-twice.yaml" with wire_server(respond) as wire: _gemini_config(path, wire.url) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway responses: Final = tuple( candidate.request( @@ -604,7 +604,7 @@ def test_gemini_passthrough_repeated_errors_each_get_row_and_log_line(gateway: G def test_budget_rejected_call_keeps_budget_normalized_error(gateway: Gateway, tmp_path: Path) -> None: path: Final = tmp_path / "budget.yaml" path.write_text(Path("tests/integration/proxy_config.yaml").read_text()) - with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: candidate: Final = owned.gateway with candidate.scenario() as scenario: model: Final = scenario.model() diff --git a/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py b/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py index 003181c72d8..d35a644ea3d 100644 --- a/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py +++ b/tests/integration/observability/test_s3_v2_streaming_failure_dedupe.py @@ -174,7 +174,6 @@ def test_retried_failure_uploads_one_s3_object( tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, config=config, - workers=2, ) as candidate, candidate.scenario() as scenario, ): @@ -216,7 +215,6 @@ def test_fallback_chain_failure_uploads_one_s3_object( tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, config=config, - workers=2, ) as candidate, candidate.scenario() as scenario, ): @@ -268,7 +266,6 @@ def test_streaming_success_uploads_one_s3_object( tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, config=config, - workers=2, ) as candidate, candidate.scenario() as scenario, ): diff --git a/tests/integration/spend/test_spend_log_tool_payload_content.py b/tests/integration/spend/test_spend_log_tool_payload_content.py index 22ed6c74756..672c0a38407 100644 --- a/tests/integration/spend/test_spend_log_tool_payload_content.py +++ b/tests/integration/spend/test_spend_log_tool_payload_content.py @@ -325,7 +325,7 @@ def test_stored_chat_response_keeps_logprob_tokens(gateway: Gateway, tmp_path: P with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model( @@ -418,7 +418,7 @@ def test_stored_messages_keep_tool_use_input(gateway: Gateway, tmp_path: Path) - with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model( @@ -650,7 +650,7 @@ def test_chat_sdk_keeps_logprob_tokens(gateway: Gateway, tmp_path: Path, client_ with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model( @@ -702,7 +702,7 @@ def test_chat_sdk_stream_include_usage_masks_request_fields(gateway: Gateway, tm with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model( @@ -747,7 +747,7 @@ def test_chat_history_keeps_string_tool_arguments_and_tool_content(gateway: Gate with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") @@ -833,7 +833,7 @@ def test_messages_sdk_keeps_tool_use_input(gateway: Gateway, tmp_path: Path, cli with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -879,7 +879,7 @@ def test_messages_stream_keeps_tool_use_input(gateway: Gateway, tmp_path: Path) with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -925,7 +925,7 @@ def test_responses_stream_keeps_function_call_arguments(gateway: Gateway, tmp_pa with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -1006,7 +1006,7 @@ def test_native_responses_keeps_logprob_tokens(gateway: Gateway, tmp_path: Path) with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") @@ -1052,7 +1052,7 @@ def test_malformed_tool_blocks_keep_only_recognized_content(gateway: Gateway, tm with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") @@ -1097,7 +1097,7 @@ def test_messages_tool_input_handles_mixed_values_and_truncation(gateway: Gatewa with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -1143,7 +1143,7 @@ def test_messages_without_auth_create_no_spend_row(gateway: Gateway, tmp_path: P with ( wire_server(respond) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, ): response: Final = isolated.client.post( "/v1/messages", @@ -1182,7 +1182,7 @@ def test_messages_upstream_error_keeps_tool_input(gateway: Gateway, tmp_path: Pa with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") @@ -1235,7 +1235,6 @@ def test_store_prompts_off_keeps_chat_and_messages_representation_equal(gateway: tmp_path, {}, config=_prompt_storage_config(tmp_path, store_prompts=False), - workers=2, ) as isolated, isolated.scenario() as scenario, ): @@ -1283,7 +1282,7 @@ def test_identical_messages_requests_have_distinct_spend_rows(gateway: Gateway, with ( wire_server(_answering_model_listing(respond)) as wire, - owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path), workers=2) as isolated, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, isolated.scenario() as scenario, ): model: Final = scenario.model(model=ANTHROPIC_MODEL, api_base=wire.url, api_key="synthetic-anthropic-key") diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index 998f81155a3..b2271a1ed47 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -11,6 +11,7 @@ import threading import time from collections.abc import Iterator from contextlib import contextmanager +from functools import reduce from pathlib import Path from types import ModuleType from typing import Final @@ -238,6 +239,45 @@ def test_only_live_readable_workers_can_satisfy_readiness( 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: + probe: Final = process_module._Probe + timeline: Final = ( + probe(0.5, listening=False, prisma=False, ready_pid=None), + probe(4.0, listening=False, prisma=True, ready_pid=None), + probe(5.0, listening=False, prisma=False, ready_pid=None), + probe(5.5, listening=False, prisma=True, ready_pid=None), + probe(6.5, listening=False, prisma=False, ready_pid=None), + probe(7.0, listening=True, prisma=False, ready_pid=None), + probe(10.0, listening=True, prisma=False, ready_pid=41), + probe(11.0, listening=True, prisma=False, ready_pid=42), + ) + phases: Final = reduce(process_module._observe, timeline, process_module.BootPhases()) + assert phases == process_module.BootPhases( + prisma_started=4.0, + prisma_finished=6.5, + listening=7.0, + first_ready=10.0, + ready_pids=frozenset((41, 42)), + ) + + +def test_boot_phases_leave_prisma_unset_when_no_prisma_child_was_seen(process_module: ModuleType) -> None: + probe: Final = process_module._Probe + timeline: Final = ( + probe(1.0, listening=False, prisma=False, ready_pid=None), + probe(2.0, listening=True, prisma=False, ready_pid=7), + ) + phases: Final = reduce(process_module._observe, timeline, process_module.BootPhases()) + assert (phases.prisma_started, phases.prisma_finished, phases.listening, phases.first_ready) == ( + None, + None, + 2.0, + 2.0, + ) + + @pytest.mark.parametrize( ("writer", "explicit_reader", "expected_reader"), ( diff --git a/tests/unit/integration_support/test_workers.py b/tests/unit/integration_support/test_workers.py index 68c43335e32..30e55a91b11 100644 --- a/tests/unit/integration_support/test_workers.py +++ b/tests/unit/integration_support/test_workers.py @@ -1,12 +1,21 @@ +import threading +from contextlib import ExitStack from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final import httpx import pytest from tests.integration._support.client import Gateway -from tests.integration._support.workers import GatewayProfile, gateway_profile, readiness_failure +from tests.integration._support.workers import ( + GatewayProfile, + WorkerStack, + gateway_profile, + readiness_failure, + start_worker_stacks, +) from tests.integration.run import worker_settings HEALTHY: Final = { @@ -82,3 +91,43 @@ def test_replica_runs_and_coverage_runs_reach_every_worker_gateway() -> None: assert profile == GatewayProfile( peer=True, replica=True, cost_map=False, coverage_file=Path("/results/coverage/data") ) + + +def _stack(name: str, closed: threading.Event) -> WorkerStack: + resources: Final = ExitStack() + resources.callback(closed.set) + return WorkerStack(resources, MappingProxyType({"INTEGRATION_RESULTS_DIR": name})) + + +def test_worker_stacks_start_together_and_return_in_worker_order(tmp_path: Path) -> None: + directories: Final = tuple(tmp_path / f"gw{index}" for index in range(3)) + closed: Final = MappingProxyType({directory: threading.Event() for directory in directories}) + together: Final = threading.Barrier(len(directories), timeout=10) + + def build(directory: Path) -> WorkerStack: + together.wait() + return _stack(directory.name, closed[directory]) + + with ExitStack() as resources: + environments: Final = start_worker_stacks(directories, build, resources, len(directories)) + assert [environment["INTEGRATION_RESULTS_DIR"] for environment in environments] == ["gw0", "gw1", "gw2"] + assert not any(event.is_set() for event in closed.values()) + assert all(event.is_set() for event in closed.values()) + + +def test_a_failed_worker_stack_stops_every_sibling_that_started(tmp_path: Path) -> None: + directories: Final = tuple(tmp_path / f"gw{index}" for index in range(3)) + closed: Final = MappingProxyType({directory: threading.Event() for directory in directories}) + together: Final = threading.Barrier(len(directories), timeout=10) + + def build(directory: Path) -> WorkerStack: + together.wait() + if directory.name == "gw1": + raise RuntimeError("gw1 gateway never became ready") + return _stack(directory.name, closed[directory]) + + with ExitStack() as resources: + with pytest.raises(RuntimeError, match="gw1 gateway never became ready"): + start_worker_stacks(directories, build, resources, len(directories)) + assert closed[directories[0]].is_set() + assert closed[directories[2]].is_set()