diff --git a/.circleci/config.yml b/.circleci/config.yml index be64a022dcf..e4333c88a64 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -160,7 +160,7 @@ commands: exit 1 ;; esac - curl -sSLf -o /tmp/rustup-init \ + curl -sSLf --retry 3 --retry-all-errors --retry-max-time 120 --connect-timeout 15 --max-time 60 -o /tmp/rustup-init \ "https://static.rust-lang.org/rustup/archive/1.28.2/${RUSTUP_TRIPLE}/rustup-init" echo "${RUSTUP_SHA256} /tmp/rustup-init" | sha256sum -c - chmod +x /tmp/rustup-init @@ -1749,6 +1749,7 @@ workflows: name: integration-management suite: management parallelism: 4 + resource_class: xlarge - integration_shard_coverage: name: integration-management-coverage suite: management @@ -1758,6 +1759,7 @@ workflows: name: integration-accounting suite: accounting parallelism: 4 + resource_class: xlarge - integration_shard_coverage: name: integration-accounting-coverage suite: accounting @@ -1766,11 +1768,12 @@ workflows: - integration_contracts: name: integration-providers suite: providers - parallelism: 8 + parallelism: 4 + resource_class: xlarge - integration_shard_coverage: name: integration-providers-coverage suite: providers - shards: 8 + shards: 4 requires: - integration-providers - integration_contracts: @@ -1797,6 +1800,7 @@ workflows: name: integration-extensions suite: extensions parallelism: 8 + resource_class: xlarge - integration_shard_coverage: name: integration-extensions-coverage suite: extensions @@ -1808,6 +1812,7 @@ workflows: suite: management mode: replica parallelism: 4 + resource_class: xlarge - integration_shard_coverage: name: integration-management-replica-coverage suite: management diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh index 28d1fba68c7..7c6bee57ba4 100644 --- a/.circleci/scripts/run_integration.sh +++ b/.circleci/scripts/run_integration.sh @@ -128,9 +128,20 @@ fi sudo iptables -L integration_only -n -v -x > "$results/egress-guard.txt" awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/egress-guard.txt" -setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \ - .venv/bin/python -m integration._support.upstream > "$results/upstream.log" 2>&1 & -upstream_pid=$! +isolated_workers=false +if [ "$mode" = parity ]; then + export INTEGRATION_WORKERS=1 +else + case "$suite" in + management|accounting) export INTEGRATION_WORKERS=2; isolated_workers=true ;; + providers|extensions) export INTEGRATION_WORKERS=4; isolated_workers=true ;; + esac +fi +if [ "$isolated_workers" = false ]; then + setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \ + .venv/bin/python -m integration._support.upstream > "$results/upstream.log" 2>&1 & + upstream_pid=$! +fi if [ "$suite" = cost ]; then export INTEGRATION_WORKERS=8 fi @@ -180,17 +191,22 @@ start_proxy() { > "$results/$log_name" 2>&1 & launched_pid=$! } -start_proxy 4000 proxy.log -proxy_pid="$launched_pid" -.venv/bin/python .circleci/scripts/wait_integration_services.py -curl --noproxy '*' -sSf -X POST "$INTEGRATION_PROXY_URL/config/update" \ - -H "Authorization: Bearer $LITELLM_MASTER_KEY" -H 'Content-Type: application/json' \ - -d '{"router_settings": {"num_retries": 0}}' > "$results/seed-router-settings.json" -if [ "$suite" = management ] || [ "$suite" = mcp ]; then +if { [ "$suite" = management ] || [ "$suite" = mcp ]; } && [ "$isolated_workers" = true ]; then export INTEGRATION_PEER_URL=http://127.0.0.1:4001 - start_proxy 4001 peer.log - peer_pid="$launched_pid" +fi +if [ "$isolated_workers" = false ]; then + start_proxy 4000 proxy.log + proxy_pid="$launched_pid" .venv/bin/python .circleci/scripts/wait_integration_services.py + curl --noproxy '*' -sSf -X POST "$INTEGRATION_PROXY_URL/config/update" \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" -H 'Content-Type: application/json' \ + -d '{"router_settings": {"num_retries": 0}}' > "$results/seed-router-settings.json" + if [ "$suite" = management ] || [ "$suite" = mcp ]; then + export INTEGRATION_PEER_URL=http://127.0.0.1:4001 + start_proxy 4001 peer.log + peer_pid="$launched_pid" + .venv/bin/python .circleci/scripts/wait_integration_services.py + fi fi if [ "$suite" = providers ]; then diff --git a/.circleci/scripts/wait_integration_services.py b/.circleci/scripts/wait_integration_services.py index 486e37cba00..0649585580f 100644 --- a/.circleci/scripts/wait_integration_services.py +++ b/.circleci/scripts/wait_integration_services.py @@ -1,42 +1,31 @@ import os -import time +from contextlib import ExitStack from typing import Final import httpx +from integration._support.client import Gateway +from integration._support.workers import wait_for_services from redis import Redis def main() -> None: primary: Final = os.environ["INTEGRATION_PROXY_URL"] peer: Final = os.environ.get("INTEGRATION_PEER_URL") - proxies: Final = (primary, peer) if peer else (primary,) - deadline: Final = time.monotonic() + 90 - headers: Final = {"Authorization": f"Bearer {os.environ['INTEGRATION_MASTER_KEY']}"} - with httpx.Client(trust_env=False, timeout=2) as client, Redis( - host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), socket_timeout=2 - ) as cache: - while True: - try: - ready: Final = ( - client.get(f"{os.environ['INTEGRATION_UPSTREAM_URL']}/health").status_code == 200 - and all(client.get(f"{url}/health/readiness").status_code == 200 for url in proxies) - ) - if ready: - for url in proxies: - response: Final = client.get(f"{url}/cache/ping", headers=headers) - response.raise_for_status() - result: Final = response.json() - assert result["status"] == "healthy", result - assert result["cache_type"] == "redis", result - assert result["ping_response"] is True, result - assert result["set_cache_response"] == "success", result - if cache.pubsub_numsub("litellm_proxy.auth_cache_invalidation")[0][1] >= len(proxies): - return - except httpx.TransportError: - pass - if time.monotonic() >= deadline: - raise SystemExit("Integration services or auth-cache subscribers did not become ready") - time.sleep(0.2) + urls: Final = (primary, peer) if peer else (primary,) + upstream: Final = os.environ["INTEGRATION_UPSTREAM_URL"] + with ExitStack() as stack: + proxies: Final = tuple( + Gateway( + stack.enter_context(httpx.Client(base_url=url, trust_env=False, timeout=2)), + os.environ["INTEGRATION_MASTER_KEY"], + upstream, + ) + for url in urls + ) + cache: Final = stack.enter_context( + Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), socket_timeout=2) + ) + wait_for_services(upstream, proxies, cache) if __name__ == "__main__": diff --git a/tests/e2e/ui/tests/migration/migratedPages.spec.ts b/tests/e2e/ui/tests/migration/migratedPages.spec.ts index f2dbd2dbc77..c5f7e5effd4 100644 --- a/tests/e2e/ui/tests/migration/migratedPages.spec.ts +++ b/tests/e2e/ui/tests/migration/migratedPages.spec.ts @@ -6,6 +6,7 @@ import { proxyIsPremium } from "../../helpers/premium"; const ROOT = (process.env.SERVER_ROOT_PATH ?? "").replace(/\/+$/, ""); const apiKeys = MIGRATED_E2E_PAGES["api-keys"]; +const home: MigratedPage = { segment: "", linkName: "Home", content: { role: "heading", name: "Home" } }; async function expectContent(page: Page, destination: MigratedPage): Promise { await expect(sidebarLink(page, apiKeys.linkName)).toBeVisible({ timeout: 20_000 }); @@ -40,7 +41,7 @@ test.describe("App Router migrated pages", () => { const landing = await page.goto(`${ROOT}/ui/`); expect(landing?.ok(), "dashboard document loads successfully").toBe(true); await dismissFeedbackPopup(page); - await expectContent(page, apiKeys); + await expectContent(page, home); await navigateToDestination(page, destination); @@ -63,7 +64,7 @@ test.describe("App Router migrated pages", () => { const landing = await page.goto(`${ROOT}/ui/`); expect(landing?.ok(), "dashboard document loads successfully").toBe(true); await dismissFeedbackPopup(page); - await expectContent(page, apiKeys); + await expectContent(page, home); for (const destination of [apiKeys, MIGRATED_E2E_PAGES.models, apiKeys]) { await navigateToDestination(page, destination); diff --git a/tests/integration/README.md b/tests/integration/README.md index d9f8275b32d..a90faa0823d 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -6,13 +6,13 @@ The ROI database contracts in `database/test_roi_observed.py` run as plain pytes The `cost` group is driven by the `CostTrackingTestCase` literals under `cost_calculation/cases/.py` (gathered by `cost_calculation/catalog.py`), with the scripted cost map in `cost_calculation/cost_map.py`. Each case has a name, contract ID, cost-map model, optional deployment overrides, request body, tagged response and exact or recount expectations. Request bodies use `$MODEL` for the registered proxy model, while responses use `$REQUEST_ID` for the per-run scenario ID. To add a case, add a cost-map entry when the model is new, add the request body and exact provider response data, and add hand-computed expected values. The upstream serves each stored response for any path under `/`, while the test-owned cost map is served over loopback through `LITELLM_MODEL_COST_MAP_URL` -Use `tests/integration/run.py management`, `accounting`, `database`, `providers`, `extensions`, `mcp`, `sdk` or `cost` to run a selected group. The group to directory mapping is the `GROUPS` literal at the top of `run.py`; a new directory needs a `GROUPS` entry and an `OWNED_DIRECTORIES` entry in `_support/manifest.py`. Set `INTEGRATION_WORKERS` above 1 to run a group under pytest-xdist; the `mcp` job does this in CI, so MCP tests must own their resources per scenario. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. When that deployment runs more than one proxy worker, set `INTEGRATION_PROXY_WORKERS` to the count so a test that writes a model and then calls it waits out the config reload interval, the only cross-worker convergence bound the wire exposes. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate +Use `tests/integration/run.py management`, `accounting`, `database`, `providers`, `extensions`, `mcp`, `sdk` or `cost` to run a selected group. The group to directory mapping is the `GROUPS` literal at the top of `run.py`; a new directory needs a `GROUPS` entry and an `OWNED_DIRECTORIES` entry in `_support/manifest.py`. Set `INTEGRATION_WORKERS` above 1 to run a group under pytest-xdist. Management, accounting, providers and extensions give each worker a private database, Redis, upstream, fake provider and proxy, with a peer and read-only database role where required. The pytest controller owns these services and closes them after worker failures. Their cases use pytest-xdist load groups inside the files selected by CircleCI timing splits. Module and class fixtures are grouped automatically, so authors do not need scheduling markers. Function fixtures remain independently schedulable. Tests that own an independent mutable sink for every case should use function fixtures, so a shared module fixture does not serialize them or share pending observations. Nested harness probes discard the parent run's shard, worker and result settings, then supply their own selections. MCP keeps its existing resource ownership and explicit load groups. Automatic group suffixes are removed from JUnit and execution records to preserve case identities. Owned proxy helpers wait for a healthy response from every configured, live worker before yielding. Serial runs keep the supplied deployment, including routing parity runs. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. When that deployment runs more than one proxy worker, set `INTEGRATION_PROXY_WORKERS` to the count so a test that writes a model and then calls it waits out the config reload interval, the only cross-worker convergence bound the wire exposes. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate Management also requires `INTEGRATION_PEER_URL`, `REDIS_HOST` and `REDIS_PORT`. CircleCI starts two directly addressed proxy processes sharing only that job's stores. The test-only CLI wrapper supplies enterprise route entitlement, following the existing behavior suite's convention. It does not qualify license validation; run it with one worker and no reload -The generated lifecycle models use 20 examples, eight steps, generation and shrinking, with isolated resources per example. HTTP operation caps include generation and shrinking and exempt cleanup. Local qualification defaults to seed 4106601 and canonical order; CircleCI derives exploration and ordering seeds from the checked-out revision and workflow ID. The ordering seed shuffles files, scoped fixture configurations and cases within each configuration. Cases sharing a module or class fixture parameter stay together, so shuffling does not repeatedly rebuild the same configuration +The generated lifecycle models use 20 examples, eight steps, generation and shrinking, with isolated resources per example. HTTP operation caps include generation and shrinking and exempt cleanup. Local qualification defaults to seed 4106601 and canonical order; CircleCI derives exploration and ordering seeds from the checked-out revision and workflow ID. The ordering seed shuffles files, scoped fixture configurations and cases within each configuration. Ordering keeps shared module and class configurations together where their dependencies permit it. Dependent fixtures with their own module parameters can still rebuild an outer configuration, with every case retained -CircleCI uses its built-in `circleci tests split --split-by=timings --timings-type=filename` to assign test files using the uploaded JUnit timings. Each file stays on one node, preserving its module and parametrized fixtures. Jobs default to one node. Parallelism should use the smallest node count that keeps the complete job, including setup and cleanup, within twenty minutes. A group that already finishes within that target on one node stays on one node. Management, accounting and management replica currently use four nodes. Providers and extensions use eight nodes; database and MCP use two, and fast suites use one. No job is configured above eight nodes. A single slow file can still determine a suite's runtime; split such a file along independent fixture boundaries rather than adding more nodes +CircleCI uses its built-in `circleci tests split --split-by=timings --timings-type=filename` to assign test files using the uploaded JUnit timings. Each file stays on one node, preserving its module and parametrized fixtures. Jobs default to one node. Parallelism should use the smallest node count that keeps the complete job, including setup and cleanup, within twenty minutes. A group that already finishes within that target on one node stays on one node. Management, accounting, providers and management replica currently use four nodes. Extensions uses eight nodes; database and MCP use two, and fast suites use one. No job is configured above eight nodes. A single slow file can still determine a suite's runtime; split such a file along independent fixture boundaries rather than adding more nodes Every node collects the same candidate files before applying CircleCI's assigned file list. Its execution record retains the candidate inventory and the selected, passed and skipped cases. An empty assigned file list produces an empty shard, not a full-suite fallback. Every parallel group has a dependent census job that rejects an empty candidate inventory, missing shards, omitted cases and duplicate execution. Standard and replica runs keep separate workspace records @@ -24,7 +24,7 @@ Reuse the existing canned provider handlers through `_support/upstream.py`. It r Owned proxies disable `.env` loading by default. A contract testing dotenv behavior can explicitly set `PYTHON_DOTENV_DISABLED=0` in its owned environment. An explicit scratch writer database drops an inherited reader with a different database name; pass both database URLs when the contract needs a paired reader. Database relays close listeners, active connections and their event loops on teardown. Upstream startup failures reap the child process before returning the failure. Custom subprocesses, direct SQL and background tasks still need an explicit owner in the test -Owned proxy workers configure a 32-thread executor before application startup, preserving additional worker startup hooks. Held synchronous provider calls must not consume a CPU-dependent executor limit before the rest of a fault burst can start. The proxy entrypoint imports LiteLLM inside `main`, so spawned workers begin supervision before importing the application. Worker healthchecks have a five-second deadline; application readiness and graceful teardown retain their separate longer deadlines. A fault that deliberately pauses Python before it can answer worker pings must supply a longer healthcheck through `extra_arguments`. Slow application startup is covered separately by the readiness deadline. Crash tests own their client tasks with `TaskGroup`, release held peers in `finally` and prove that the intended worker had active traffic before killing it +Owned proxy workers configure a 32-thread executor before application startup, preserving additional worker startup hooks. Held synchronous provider calls must not consume a CPU-dependent executor limit before the rest of a fault burst can start. The proxy entrypoint imports LiteLLM inside `main`, so spawned workers begin supervision before importing the application. Worker healthchecks have a five-second deadline; application readiness and graceful teardown retain their separate longer deadlines. A fault that deliberately pauses Python before it can answer worker pings must supply a longer healthcheck through `extra_arguments`. Slow application startup is covered separately by the readiness deadline. Cases with two complete owned proxy lifetimes use a 180-second pytest budget for startup, assertions and teardown; their HTTP, provider and persistence deadlines stay unchanged. Readiness requires responsive worker identities that are still alive and a positively identified owned listener. Unreadable socket records cannot establish ownership. Crash tests own their client tasks with `TaskGroup`, release held peers in `finally` and verify that the intended worker dies. A test requiring traffic on both workers prepares that traffic explicitly instead of assuming OS load balancing The CircleCI workflow starts its own database and Redis, restricts test-phase egress to its owned services and writes JUnit plus an executed-node manifest. Missing setup, failed cleanup or a selected test with neither a passed call nor a skip fail qualification. Skipped nodes are listed under `skipped` in `execution.json`, so the skip reasons double as the open bug list. `GITHUB_FILES` in `run.py` lists the files excluded from the `run_integration.sh` selection @@ -42,7 +42,7 @@ Accounting cases compare persisted input and output cost components against lite Provider contracts exercise actual TCP requests with synthetic credentials and local protocol peers. The S3 verifier uses independently implemented equations, a published known-answer vector, a fixed signing clock and deliberately invalid signed requests. Bedrock cases clear ambient AWS credential sources and check the literal model path, loaded role references, STS requests and bearer-only behavior -CircleCI runs providers on eight isolated nodes using its file timing splitter. Each node owns its services. The dependent providers coverage job verifies all eight execution records against the candidate case inventory and rejects missing, duplicate, failed or incomplete execution. The node's assigned filenames are saved in `node-files.txt` beside its execution record +CircleCI runs providers on four isolated nodes using its file timing splitter. Each node owns its services. The dependent providers coverage job verifies all four execution records against the candidate case inventory and rejects missing, duplicate, failed or incomplete execution. The node's assigned filenames are saved in `node-files.txt` beside its execution record Streaming checks send real HTTP transfer chunks, including one-byte partitions, fragmented tools, incomplete transfers and a cancellation barrier. They assert meaningful text, tool arguments, final usage and persisted cost. The Redis recovery case owns a separate database and Redis process, uses the supported one-second circuit-breaker recovery setting, waits for the real subscriber and verifies response data in Redis after restart. CircleCI reuses its existing Redis image for that extra process; it never pulls an image during tests @@ -50,7 +50,7 @@ The `messages_endpoint/` directory holds `/v1/messages` endpoint contracts: nati The sdk shard exercises the SDK's own HTTP clients against local protocol peers with no gateway in the path, so a case here fails only when the client library or its wire behavior changes. The HTTP/2 case runs a hypercorn TLS peer offering h2 and http/1.1 over ALPN, drives the sync and async httpx handlers at it with `LITELLM_HTTP2` off and on, and asserts the version both the client and the peer observed on the wire. Put a test here only when it needs no proxy or database. CircleCI starts a local Redis for this shard like the others, so SDK-side caching cases that need a real Redis server belong here too; a case that reaches the gateway belongs in one of the other shards -The extensions shard uses the built-in generic callback and guardrail transports. It checks callback correlation and credential exclusion, guardrail rewriting and denial, retained OpenAI consumers and A2A wire versions. CircleCI runs it on eight parallel large nodes using its file timing splitter. Each node starts its own database, Redis, upstream and proxy and runs its assigned files serially. The dependent coverage job verifies all eight execution records against the candidate case inventory. Tests must not assume a particular set of sibling files +The extensions shard uses the built-in generic callback and guardrail transports. It checks callback correlation and credential exclusion, guardrail rewriting and denial, retained OpenAI consumers and A2A wire versions. CircleCI runs it on eight xlarge nodes using its file timing splitter. Each node runs its assigned files with four pytest workers, each owning its database, Redis, upstream and proxy. The dependent coverage job verifies all eight execution records against the candidate case inventory. Tests must not assume a particular set of sibling files The security shard uses four pytest workers on an xlarge machine, each with a private database and Redis. Route sweeps use a new HTTP client for every call and share only the TLS verification context, so authentication and response cookies cannot cross callers. The canary rig supplies owned OpenAI and Anthropic endpoints and synthetic credentials and rejects attempted external egress diff --git a/tests/integration/_support/asgi.py b/tests/integration/_support/asgi.py index eff01a6ab13..e84dc9439d3 100644 --- a/tests/integration/_support/asgi.py +++ b/tests/integration/_support/asgi.py @@ -14,18 +14,20 @@ from starlette.types import ASGIApp @contextmanager -def asgi_server(app: ASGIApp, *, before_stop: Callable[[], None] | None = None) -> Iterator[str]: +def asgi_server( + app: ASGIApp, *, before_stop: Callable[[], None] | None = None, port: int = 0, keepalive_seconds: int = 1 +) -> Iterator[str]: with socket.socket() as listener: - listener.bind(("127.0.0.1", 0)) - port: Final = listener.getsockname()[1] + listener.bind(("127.0.0.1", port)) + bound_port: Final = listener.getsockname()[1] server: Final = uvicorn.Server( uvicorn.Config( app, host="127.0.0.1", - port=port, + port=bound_port, lifespan="on", log_level="warning", - timeout_keep_alive=1, + timeout_keep_alive=keepalive_seconds, timeout_graceful_shutdown=5, ) ) @@ -58,7 +60,7 @@ def asgi_server(app: ASGIApp, *, before_stop: Callable[[], None] | None = None) while not server.started: assert worker.is_alive() and time.monotonic() < deadline, "Owned ASGI peer failed readiness" time.sleep(0.01) - yield f"http://127.0.0.1:{port}" + yield f"http://127.0.0.1:{bound_port}" finally: if before_stop is not None: before_stop() diff --git a/tests/integration/_support/ordering.py b/tests/integration/_support/ordering.py index fc10c372a93..29617be62b9 100644 --- a/tests/integration/_support/ordering.py +++ b/tests/integration/_support/ordering.py @@ -24,13 +24,17 @@ class CollectedCase: fixtures: tuple[str, ...] = () +def _fixture_owner(item: pytest.Function, scope: str, baseid: str) -> str: + parent: Final = item.getparent(pytest.Class) if scope == "class" else None + return parent.nodeid if parent is not None else baseid + + def _parameter(item: pytest.Function, name: str, value: object) -> ScopedParameter | None: definitions: Final = item._fixtureinfo.name2fixturedefs.get(name) scope: Final = item.callspec._arg2scope[name].value if not definitions or scope == "function" or definitions[-1].func is get_direct_param_fixture_func: return None - parent: Final = item.getparent(pytest.Class) if scope == "class" else None - owner: Final = parent.nodeid if parent is not None else definitions[-1].baseid + owner: Final = _fixture_owner(item, scope, definitions[-1].baseid) return ScopedParameter(f"{scope}:{owner}", name, value) @@ -38,7 +42,7 @@ def collected_case(item: pytest.Item) -> CollectedCase: if not isinstance(item, pytest.Function): return CollectedCase(item.nodeid) fixtures: Final = tuple( - f"{definitions[-1].scope}:{definitions[-1].baseid}:{name}" + f"{definitions[-1].scope}:{_fixture_owner(item, definitions[-1].scope, definitions[-1].baseid)}:{name}" for name, definitions in item._fixtureinfo.name2fixturedefs.items() if definitions and definitions[-1].scope in ("module", "class") @@ -70,6 +74,14 @@ def _same_parameters(first: CollectedCase, second: CollectedCase) -> bool: return False if not first.parameters and first.fixtures != second.fixtures: return False + first_owners: Final = frozenset( + fixture.rsplit(":", 1)[0] for fixture in first.fixtures if fixture.startswith("class:") + ) + second_owners: Final = frozenset( + fixture.rsplit(":", 1)[0] for fixture in second.fixtures if fixture.startswith("class:") + ) + if first_owners != second_owners: + return False return all( left.owner == right.owner and left.name == right.name and _equal(left.value, right.value) for left, right in zip(first.parameters, second.parameters, strict=True) @@ -93,3 +105,19 @@ def _file_groups(cases: Iterable[CollectedCase]) -> tuple[tuple[str, str], ...]: def fixture_groups(cases: tuple[CollectedCase, ...]) -> Mapping[str, str]: files: Final = groupby(sorted(cases, key=_file), key=_file) return MappingProxyType(dict(chain.from_iterable(_file_groups(members) for _, members in files))) + + +def _scheduling_case(case: CollectedCase) -> CollectedCase: + if not any(fixture.startswith("module:") for fixture in case.fixtures) and not any( + parameter.owner.startswith("module:") for parameter in case.parameters + ): + return case + return CollectedCase( + case.nodeid, + tuple(parameter for parameter in case.parameters if not parameter.owner.startswith("class:")), + tuple(fixture for fixture in case.fixtures if not fixture.startswith("class:")), + ) + + +def scheduling_groups(cases: tuple[CollectedCase, ...]) -> Mapping[str, str]: + return fixture_groups(tuple(_scheduling_case(case) for case in cases)) diff --git a/tests/integration/_support/otlp_sink.py b/tests/integration/_support/otlp_sink.py index a4d2ea30959..9c5f0150b42 100644 --- a/tests/integration/_support/otlp_sink.py +++ b/tests/integration/_support/otlp_sink.py @@ -19,16 +19,19 @@ import subprocess import sys import threading import time -from collections.abc import Iterator, Mapping, Sequence +import uuid +from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import ExitStack, contextmanager from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path +from types import MappingProxyType from typing import Final from urllib.parse import urlparse import httpx import psutil +import yaml from pydantic import JsonValue, TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -364,6 +367,47 @@ class SpanSinks: arize: str +AuditConfigWriter = Callable[[Path, Mapping[str, JsonValue]], Path] +_AUDIT_CONFIGURATION: Final = TypeAdapter(dict[str, JsonValue]) + + +def audit_configuration_writer(sinks: SpanSinks) -> AuditConfigWriter: + tenant_host: Final = urlparse(sinks.tenant).netloc + + def write(directory: Path, litellm_settings: Mapping[str, JsonValue] = MappingProxyType({})) -> Path: + config: Final = _AUDIT_CONFIGURATION.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + settings: Final = _AUDIT_CONFIGURATION.validate_python(config.get("litellm_settings", {})) + general: Final = _AUDIT_CONFIGURATION.validate_python(config.get("general_settings", {})) + configured: Final = { + **config, + "litellm_settings": { + **settings, + "callbacks": ["otel"], + "provider_url_destination_allowed_hosts": [tenant_host], + **litellm_settings, + }, + "callback_settings": { + "otel": {"exporter": "http/json", "endpoint": sinks.operator, "use_simple_processor": True} + }, + "general_settings": {**general, "user_api_key_cache_ttl": 2}, + } + path: Final = directory / f"otel-audit-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(configured)) + return path + + return write + + +def langfuse_callback_variables(sinks: SpanSinks) -> dict[str, JsonValue]: + return { + "langfuse_public_key": "pk-lf-audit", + "langfuse_secret_key": "sk-lf-audit", + "langfuse_host": sinks.tenant, + } + + @dataclass(frozen=True, slots=True) class GrpcSink: url: str diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index c4c037603ff..4f8ca036489 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -1,6 +1,7 @@ import errno import json import os +import re import signal import socket import subprocess @@ -12,6 +13,7 @@ from contextlib import contextmanager from dataclasses import dataclass from datetime import UTC, datetime, timedelta from ipaddress import IPv4Address +from itertools import accumulate from pathlib import Path from types import MappingProxyType from typing import Final @@ -19,15 +21,19 @@ from urllib.parse import urlsplit import httpx import psutil +import yaml from cryptography import x509 from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.x509.oid import NameOID -from integration._support.client import GATEWAY_LIMITS, Gateway +from pydantic import TypeAdapter + +from tests.integration._support.client import GATEWAY_LIMITS, Gateway DB_PUSH: Final = ("--use_prisma_db_push",) MIGRATE_DEPLOY: Final = () LEGACY_MIGRATE_DEPLOY: Final = ("--use_legacy_migration_resolver",) +CONFIGURATION: Final = TypeAdapter(dict[str, object]) def proxy_database_environment() -> Mapping[str, str]: @@ -123,6 +129,7 @@ def _stop(process: subprocess.Popen[bytes]) -> None: signal_group(process.pid, signal.SIGKILL) process.wait(timeout=3) psutil.wait_procs(group_members(process.pid), timeout=3) + assert not group_members(process.pid), "Owned proxy child survived emergency cleanup" raise @@ -189,25 +196,70 @@ def _lost_port_race(exit_code: int | None, log: Path) -> bool: return exit_code is not None and _BIND_COLLISION in log.read_text().lower() -def _wait_until_ready(launch: _Launch) -> None: - with httpx.Client(base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False) as client: +def _owns_listener(launch: _Launch) -> bool: + for member in group_members(launch.process.pid): + try: + if any( + connection.status == psutil.CONN_LISTEN and connection.laddr.port == launch.port + for connection in member.net_connections(kind="inet") + ): + return True + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + return False + + +def _readiness_pids(launch: _Launch) -> Iterator[int]: + with httpx.Client( + base_url=f"http://127.0.0.1:{launch.port}", + timeout=15, + trust_env=False, + limits=httpx.Limits(max_keepalive_connections=0), + ) as client: deadline: Final = time.monotonic() + float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS", "70")) while launch.process.poll() is None: try: - if client.get("/health/readiness", timeout=2).status_code == 200: - return + 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 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 _ready_worker_pid(member: psutil.Process) -> int | None: + try: + status: Final = member.status() + return ( + None + if status in (psutil.STATUS_ZOMBIE, psutil.STATUS_DEAD, psutil.STATUS_STOPPED, psutil.STATUS_TRACING_STOP) + else member.pid + ) + except (psutil.NoSuchProcess, psutil.AccessDenied): + return None + + +def _wait_until_ready(launch: _Launch, workers: int = 1) -> None: + for ready in accumulate(_readiness_pids(launch), _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 _launch_until_bound( - command: tuple[str, ...], root: Path, environment: Mapping[str, str], output: Path, attempts: int + command: tuple[str, ...], root: Path, environment: Mapping[str, str], output: Path, attempts: int, workers: int ) -> _Launch: launch: Final = _launch(command, root, environment, output) try: - _wait_until_ready(launch) + _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" @@ -228,13 +280,42 @@ def _launch_until_bound( ) return launch _stop_launch(launch) - return _launch_until_bound(command, root, environment, output, attempts - 1) + return _launch_until_bound(command, root, environment, output, attempts - 1, workers) 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": + 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']}" + + 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) + destination: Final = directory / f"worker-config-{uuid.uuid4().hex}.yaml" + loaded: Final[object] = yaml.safe_load(contents) + if not isinstance(loaded, dict): + destination.write_text(contents) + return destination + settings: Final = CONFIGURATION.validate_python(loaded) + general: Final = settings.get("general_settings", {}) + if not isinstance(general, dict): + destination.write_text(contents) + return destination + configured: Final = { + **settings, + "general_settings": {"database_connection_pool_limit": 2, **CONFIGURATION.validate_python(general)}, + } + destination.write_text(yaml.safe_dump(configured, sort_keys=False)) + return destination + + def _proxy_environment( gateway: Gateway, overrides: Mapping[str, str], remove_environment: tuple[str, ...] ) -> Mapping[str, str]: @@ -261,7 +342,13 @@ def _proxy_environment( "PYTHON_DOTENV_DISABLED": "1", **overrides, "LITELLM_WORKER_STARTUP_HOOKS": ",".join( - hook for hook in ("integration._support.runtime:configure_executor", configured_hooks) if hook + hook + for hook in ( + "integration._support.runtime:configure_executor", + "integration._support.runtime:configure_readiness", + configured_hooks, + ) + if hook ), } ) @@ -314,7 +401,7 @@ def owned_proxy_process( "-m", "integration._support.proxy", "--config", - str(config or "tests/integration/proxy_config.yaml"), + str(worker_configuration(config or root / "tests/integration/proxy_config.yaml", directory)), "--host", "127.0.0.1", "--num_workers", @@ -324,7 +411,7 @@ def owned_proxy_process( *database_setup, *extra_arguments, ) - launch: Final = _launch_until_bound(command, root, environment, output, _PORT_ATTEMPTS) + launch: Final = _launch_until_bound(command, root, environment, output, _PORT_ATTEMPTS, workers) process: Final = launch.process try: with httpx.Client( @@ -358,7 +445,7 @@ def owned_gateway_image( "--timeout-worker-healthcheck", str(int(graceful_stop_seconds())), ) - launch: Final = _launch_until_bound(command, root, environment, output, _PORT_ATTEMPTS) + launch: Final = _launch_until_bound(command, root, environment, output, _PORT_ATTEMPTS, workers) try: with httpx.Client( base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False, limits=GATEWAY_LIMITS diff --git a/tests/integration/_support/provider.py b/tests/integration/_support/provider.py index 80c14564e25..fa0bb7ad379 100644 --- a/tests/integration/_support/provider.py +++ b/tests/integration/_support/provider.py @@ -8,6 +8,7 @@ The proxy's model-info refresh is answered here. from __future__ import annotations +import os from collections import deque from collections.abc import Iterator from contextlib import contextmanager @@ -16,7 +17,7 @@ from typing import Final from tests.integration._support.wire import Reply, Request, Wire, wire_server -PROVIDER_PORT: Final = 8191 +PROVIDER_PORT: Final = int(os.environ.get("INTEGRATION_PROVIDER_PORT", "8191")) PROVIDER_URL: Final = f"http://127.0.0.1:{PROVIDER_PORT}" MODEL_DISCOVERY: Final = ("GET", "/v1/models") _UNQUEUED: Final = Reply(status=500, body=b'{"error": "the shared fake provider has no reply queued for this request"}') @@ -32,9 +33,7 @@ class SharedProvider: self.replies.extend(replies) def received(self) -> tuple[Request, ...]: - return tuple( - request for request in self.wire.drain() if (request.method, request.target) != MODEL_DISCOVERY - ) + return tuple(request for request in self.wire.drain() if (request.method, request.target) != MODEL_DISCOVERY) @contextmanager diff --git a/tests/integration/_support/runtime.py b/tests/integration/_support/runtime.py index a709ab3d70d..a638dcfa4bb 100644 --- a/tests/integration/_support/runtime.py +++ b/tests/integration/_support/runtime.py @@ -1,8 +1,32 @@ import asyncio +import os from concurrent.futures import ThreadPoolExecutor +from typing import Final + +from fastapi import Response +from fastapi.responses import JSONResponse +from pydantic import TypeAdapter + +READINESS_BODY: Final = TypeAdapter(dict[str, object]) def configure_executor() -> None: asyncio.get_running_loop().set_default_executor( ThreadPoolExecutor(max_workers=32, thread_name_prefix="integration-provider") ) + + +async def worker_readiness() -> JSONResponse: + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + response: Final = Response() + body: Final = READINESS_BODY.validate_python(await health_readiness(response)) + return JSONResponse( + content=body, status_code=response.status_code, 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) diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index 6cac8f7276f..6b36d56a757 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -42,7 +42,7 @@ from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route, WebSocketRoute -from starlette.websockets import WebSocket +from starlette.websockets import WebSocket, WebSocketDisconnect JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) INTERNAL_FIELDS: Final = frozenset( @@ -405,6 +405,14 @@ class Provider: if not isinstance(response, RealtimeResponse): await websocket.close(code=4404) return + try: + await self._realtime_frames(websocket, response, scenario_id, authorization, api_key) + except WebSocketDisconnect: + return + + async def _realtime_frames( + self, websocket: WebSocket, response: RealtimeResponse, scenario_id: str, authorization: str, api_key: str + ) -> None: await websocket.accept() session: Final = _realtime_session(response, scenario_id, websocket.query_params.get("model", "")) for _ in range(response.created_repeats): diff --git a/tests/integration/_support/workers.py b/tests/integration/_support/workers.py index 85350caac9a..2502e62477e 100644 --- a/tests/integration/_support/workers.py +++ b/tests/integration/_support/workers.py @@ -1,16 +1,123 @@ from __future__ import annotations +import os +import socket +import time +import uuid from collections.abc import Iterator, Mapping -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from pathlib import Path from types import MappingProxyType +from typing import Final, Protocol +from urllib.parse import urlsplit, urlunsplit +import httpx +import psycopg +from psycopg import sql +from pydantic import TypeAdapter +from redis import Redis +from redis.exceptions import ConnectionError as RedisConnectionError +from redis.exceptions import TimeoutError as RedisTimeoutError + +from tests.integration._support.client import Gateway from tests.integration._support.database import scratch_database from tests.integration._support.redis_process import owned_redis +READINESS_OBJECT: Final = TypeAdapter(dict[str, object]) +SUBSCRIBER_COUNTS: Final = TypeAdapter(tuple[tuple[str, int], ...]) + + +class ServiceSubscribers(Protocol): + def pubsub_numsub(self, *channels: str) -> object: ... + + +def readiness_failure( + client: httpx.Client, upstream: str, proxies: tuple[Gateway, ...], cache: ServiceSubscribers +) -> str | None: + try: + if client.get(upstream + "/health", timeout=2).status_code != 200: + return "scripted upstream health" + for index, proxy in enumerate(proxies): + if proxy.client.get("/health/readiness", timeout=2).status_code != 200: + return f"proxy {index} readiness" + response: Final = proxy.client.get( + "/cache/ping", headers={"Authorization": f"Bearer {proxy.key}"}, timeout=2 + ) + if response.status_code != 200: + return f"proxy {index} Redis health HTTP {response.status_code}" + result: Final = READINESS_OBJECT.validate_json(response.content) + if ( + result.get("status") != "healthy" + or result.get("cache_type") != "redis" + or result.get("ping_response") is not True + or result.get("set_cache_response") != "success" + ): + return f"proxy {index} Redis read/write health" + subscribers: Final = SUBSCRIBER_COUNTS.validate_python( + cache.pubsub_numsub("litellm_proxy.auth_cache_invalidation") + ) + if subscribers[0][1] < len(proxies): + return f"auth-cache subscribers {subscribers[0][1]}/{len(proxies)}" + except (httpx.TransportError, RedisConnectionError, RedisTimeoutError) as error: + return f"service transport {type(error).__name__}" + return None + + +def wait_for_services(upstream: str, proxies: tuple[Gateway, ...], cache: Redis) -> None: + with httpx.Client(trust_env=False, timeout=2) as client: + deadline: Final = time.monotonic() + 90 + while True: + failure: Final = readiness_failure(client, upstream, proxies, cache) + if failure is None: + return + assert time.monotonic() < deadline, f"Integration services did not become ready: {failure}" + time.sleep(0.2) + + +def worker_endpoints() -> Mapping[str, str]: + with socket.socket() as upstream, socket.socket() as provider: + upstream.bind(("127.0.0.1", 0)) + provider.bind(("127.0.0.1", 0)) + return MappingProxyType( + { + "INTEGRATION_UPSTREAM_URL": f"http://127.0.0.1:{upstream.getsockname()[1]}", + "INTEGRATION_PROVIDER_PORT": str(provider.getsockname()[1]), + } + ) + + +def _replica_urls(database: str) -> Mapping[str, str]: + writer: Final = urlsplit(os.environ["INTEGRATION_PROXY_DATABASE_URL"]) + reader: Final = urlsplit(os.environ["INTEGRATION_PROXY_READ_REPLICA_URL"]) + database_path: Final = urlsplit(database).path + assert writer.username is not None and reader.username is not None + with psycopg.connect(database, autocommit=True) as admin: + admin.execute( + sql.SQL("ALTER DATABASE {} OWNER TO {}").format( + sql.Identifier(database_path.removeprefix("/")), sql.Identifier(writer.username) + ) + ) + admin.execute( + sql.SQL("GRANT CONNECT ON DATABASE {} TO {}").format( + sql.Identifier(database_path.removeprefix("/")), sql.Identifier(reader.username) + ) + ) + admin.execute(sql.SQL("GRANT USAGE ON SCHEMA public TO {}").format(sql.Identifier(reader.username))) + admin.execute( + sql.SQL("ALTER DEFAULT PRIVILEGES FOR ROLE {} IN SCHEMA public GRANT SELECT ON TABLES TO {}").format( + sql.Identifier(writer.username), sql.Identifier(reader.username) + ) + ) + return MappingProxyType( + { + "INTEGRATION_PROXY_DATABASE_URL": urlunsplit(writer._replace(path=database_path)), + "INTEGRATION_PROXY_READ_REPLICA_URL": urlunsplit(reader._replace(path=database_path)), + } + ) + @contextmanager -def worker_services(directory: Path) -> Iterator[Mapping[str, str]]: +def worker_services(directory: Path, *, replica: bool = False) -> Iterator[Mapping[str, str]]: directory.mkdir(parents=True, exist_ok=True) with ( scratch_database() as database, @@ -23,5 +130,42 @@ def worker_services(directory: Path) -> Iterator[Mapping[str, str]]: "INTEGRATION_PROXY_READ_REPLICA_URL": "", "REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port), + **(_replica_urls(database) if replica else {}), } ) + + +@contextmanager +def worker_gateway(directory: Path, *, peer: bool, replica: bool) -> Iterator[Mapping[str, str]]: + from tests.integration._support.asgi import asgi_server + from tests.integration._support.process import owned_proxy + from tests.integration._support.upstream import Provider + + port: Final = urlsplit(os.environ["INTEGRATION_UPSTREAM_URL"]).port + assert port is not None + key: Final = f"sk-{uuid.uuid4().hex}" + with worker_services(directory, replica=replica) as services: + overrides: Final = { + **services, + "DATABASE_URL": services["INTEGRATION_PROXY_DATABASE_URL"], + **({"DATABASE_URL_READ_REPLICA": services["INTEGRATION_PROXY_READ_REPLICA_URL"]} if replica else {}), + } + with ( + asgi_server(Provider().app(), port=port, keepalive_seconds=125) as upstream, + httpx.Client(base_url="http://127.0.0.1:1", trust_env=False) as client, + owned_proxy(Gateway(client, key, upstream), directory, overrides) as primary, + owned_proxy(primary, directory, overrides) if peer else nullcontext(None) as secondary, + ): + with Redis(host=services["REDIS_HOST"], port=int(services["REDIS_PORT"]), socket_timeout=2) as cache: + wait_for_services(upstream, (primary, secondary) if secondary is not None else (primary,), cache) + primary.post("/config/update", {"router_settings": {"num_retries": 0}}) + yield MappingProxyType( + { + **services, + "INTEGRATION_PROXY_URL": str(primary.client.base_url), + "INTEGRATION_PEER_URL": str(secondary.client.base_url) if secondary is not None else "", + "INTEGRATION_MASTER_KEY": key, + "LITELLM_MASTER_KEY": key, + "INTEGRATION_UPSTREAM_URL": upstream, + } + ) diff --git a/tests/integration/compatibility/test_integration_worker_cleanup.py b/tests/integration/compatibility/test_integration_worker_cleanup.py new file mode 100644 index 00000000000..a3d849ffe98 --- /dev/null +++ b/tests/integration/compatibility/test_integration_worker_cleanup.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path +from typing import Final + +import psutil +import psycopg +import pytest +from psycopg import sql +from pydantic import BaseModel, ConfigDict + +from tests.integration._support.database import read_rows + +FAULT_CASE: Final = """import json, os, pathlib, psutil +from urllib.parse import urlsplit + +def test_worker_hard_exit(): + worker = psutil.Process() + owner = psutil.Process(worker.ppid()) + pathlib.Path(os.environ["WORKER_CRASH_RECORD"]).write_text(json.dumps({ + "database": urlsplit(os.environ["DATABASE_URL"]).path.removeprefix("/"), + "processes": [[p.pid, p.create_time()] for p in owner.children(recursive=True) if p.pid != worker.pid], + })) + os._exit(23) +""" + + +class CrashResources(BaseModel): + model_config = ConfigDict(frozen=True) + database: str + processes: tuple[tuple[int, float], ...] + + +def _survivor(identity: tuple[int, float]) -> psutil.Process | None: + pid, created = identity + try: + process: Final = psutil.Process(pid) + return process if process.create_time() == created and process.status() != psutil.STATUS_ZOMBIE else None + except psutil.NoSuchProcess: + return None + + +def _clean_fault_resources(resources: CrashResources) -> None: + processes: Final = tuple(process for identity in resources.processes if (process := _survivor(identity))) + for process in processes: + process.terminate() + _stopped, alive = psutil.wait_procs(processes, timeout=5) + for process in alive: + process.kill() + psutil.wait_procs(alive, timeout=3) + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: + admin.execute(sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(sql.Identifier(resources.database))) + + +def test_worker_crash_fails_the_run_and_releases_its_database_and_processes(tmp_path: Path) -> None: + root: Final = Path(__file__).resolve().parents[3] + case: Final = tmp_path / "test_worker_fault.py" + record: Final = tmp_path / "resources.json" + case.write_text(FAULT_CASE) + environment: Final = { + **os.environ, + "PYTHONPATH": f"{root}:{root / 'tests'}:{root / 'tests' / 'e2e'}", + "PYTEST_ADDOPTS": "", + "INTEGRATION_WORKER_ISOLATION": "1", + "INTEGRATION_PEER_URL": "", + "INTEGRATION_PROXY_READ_REPLICA_URL": "", + "INTEGRATION_RESULTS_DIR": str(tmp_path / "results"), + "WORKER_CRASH_RECORD": str(record), + } + try: + with (tmp_path / "fault-run.log").open("w") as log, pytest.raises(subprocess.CalledProcessError) as refused: + subprocess.run( + ( + sys.executable, + "-P", + "-m", + "pytest", + "-p", + "tests.integration.conftest", + "-p", + "no:cacheprovider", + "-n", + "1", + "--max-worker-restart=0", + str(case), + ), + cwd=tmp_path, + env=environment, + stdout=log, + stderr=subprocess.STDOUT, + timeout=180, + check=True, + ) + resources: Final = CrashResources.model_validate_json(record.read_text()) + assert resources.database.startswith("integration_") + assert refused.value.returncode == 1 + assert resources.processes + assert not read_rows("SELECT datname FROM pg_database WHERE datname=%s", (resources.database,)) + assert all(_survivor(identity) is None for identity in resources.processes) + finally: + if record.exists(): + recorded: Final = CrashResources.model_validate_json(record.read_text()) + assert recorded.database.startswith("integration_") + _clean_fault_resources(recorded) diff --git a/tests/integration/configuration/test_lazy_routes_flag.py b/tests/integration/configuration/test_lazy_routes_flag.py index 6135deca1c8..28930ec29b4 100644 --- a/tests/integration/configuration/test_lazy_routes_flag.py +++ b/tests/integration/configuration/test_lazy_routes_flag.py @@ -582,6 +582,12 @@ def _assert_config_declared_features_serve(owned: OwnedProxy, features: _Configu ] assert [(params["name"], params["arguments"]) for params in reached_peer] == [("add", {"marker": marker})] + bare: Final = owned.gateway.request("GET", "/langfuse") + assert bare.status_code == 200, bare.text + assert JSON.validate_json(bare.content) == {"status": "OK"}, bare.text + reached_bare: Final = tuple(request.target for request in features.langfuse.drain()) + assert reached_bare == (SELF_HOSTED_LANGFUSE,), reached_bare + forwarded: Final = owned.gateway.request("GET", "/langfuse/api/public/health") assert forwarded.status_code == 200, forwarded.text reached_langfuse: Final = tuple(request.target for request in features.langfuse.drain()) @@ -590,6 +596,13 @@ def _assert_config_declared_features_serve(owned: OwnedProxy, features: _Configu ) +def _canonical_langfuse_paths(paths: tuple[str, ...]) -> tuple[str, ...]: + configured: Final = ("/langfuse", "/langfuse/{subpath:path}") + remaining: Final = tuple(path for path in paths if path not in configured) + langfuse: Final = tuple(sorted(path for path in paths if path in configured)) + return (*remaining, *langfuse) + + def test_disable_lazy_routes_flag_serves_config_declared_features_like_the_warmed_lazy_proxy( gateway: Gateway, tmp_path: Path, provider: Wire ) -> None: @@ -601,5 +614,5 @@ def test_disable_lazy_routes_flag_serves_config_declared_features_like_the_warme warmed_openapi: Final = _openapi_paths(lazy.gateway) with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}, config=features.config) as eager: _assert_config_declared_features_serve(eager, features, provider) - assert _paths(eager.gateway) == warmed_paths + assert _canonical_langfuse_paths(_paths(eager.gateway)) == _canonical_langfuse_paths(warmed_paths) assert _openapi_paths(eager.gateway) == warmed_openapi diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 280a828e1d4..dab7c1bbd1d 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -3,22 +3,27 @@ from __future__ import annotations import hashlib import json import os -from collections.abc import Iterator, Mapping, Sequence +import re +from collections.abc import Generator, Iterator, Mapping, Sequence +from contextlib import ExitStack +from copy import copy from importlib.metadata import version from pathlib import Path from typing import Final, Protocol, runtime_checkable import httpx import pytest +from execnet import XSpec from pydantic import TypeAdapter from redis import Redis from tests.integration._support.client import Gateway, eventually, gateway_from_environment from tests.integration._support.generation import LIFECYCLE_SETTINGS from tests.integration._support.manifest import OWNED_DIRECTORIES -from tests.integration._support.ordering import collected_case, fixture_groups +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 worker_endpoints, worker_gateway from tests.integration.run import GITHUB_FILES COLLECTED: Final = pytest.StashKey[tuple[str, ...]]() @@ -27,7 +32,9 @@ REPORTS: Final = pytest.StashKey[list[pytest.TestReport]]() WORKER_INVENTORY: Final = TypeAdapter(tuple[str, ...]) SHARD_PATH: Final[TypeAdapter[Path | None]] = TypeAdapter(Path | None) SHARD_FILES: Final = pytest.StashKey[frozenset[str] | None]() +GROUP_SUFFIX: Final = re.compile(r"@integration_fixture_[0-9a-f]{64}$") INTEGER_OPTION: Final = TypeAdapter(int) +STRING_ENVIRONMENT: Final = TypeAdapter(dict[str, str]) @runtime_checkable @@ -62,6 +69,37 @@ def pytest_configure(config: pytest.Config) -> None: config.pluginmanager.register(RoutingPlugin(config)) +@pytest.hookimpl(optionalhook=True) +def pytest_xdist_setupnodes(config: pytest.Config, specs: Sequence[XSpec]) -> None: + if ( + os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1" + or config.getoption("collectonly") is True + or config.stash[SHARD_FILES] == frozenset() + ): + return + resources: Final = ExitStack() + config.add_cleanup(resources.close) + destination: Final = Path(os.environ["INTEGRATION_RESULTS_DIR"]) + peer_needed: Final = bool(os.environ.get("INTEGRATION_PEER_URL")) + replica_needed: Final = bool(os.environ.get("INTEGRATION_PROXY_READ_REPLICA_URL")) + 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", peer=peer_needed, replica=replica_needed) + ) + spec.env = { + **STRING_ENVIRONMENT.validate_python(spec.env), + **endpoints, + **services, + "INTEGRATION_RESULTS_DIR": str(results), + } + + class IntegrationReportPlugin: def __init__(self, config: pytest.Config) -> None: self.config = config @@ -71,7 +109,7 @@ class IntegrationReportPlugin: @pytest.hookimpl(optionalhook=True) def pytest_xdist_node_collection_finished(self, node: object, ids: Sequence[str]) -> None: - self.config.stash[COLLECTED] = tuple(nodeid for nodeid in ids if _owned(nodeid)) + self.config.stash[COLLECTED] = tuple(_case_identity(nodeid) for nodeid in ids if _owned(nodeid)) @pytest.hookimpl(optionalhook=True) def pytest_testnodedown(self, node: WorkerNode, error: object) -> None: @@ -89,6 +127,29 @@ class IntegrationReportPlugin: output["integration_inventory"] = self.config.stash.get(INVENTORY, ()) +def _case_identity(nodeid: str) -> str: + return GROUP_SUFFIX.sub("", nodeid) + + +@pytest.hookimpl(wrapper=True, optionalhook=True) +def pytest_report_from_serializable() -> Generator[ + None, pytest.TestReport | pytest.CollectReport | None, pytest.TestReport | pytest.CollectReport | None +]: + report: Final = yield + if not isinstance(report, pytest.TestReport): + return report + normalized: Final = copy(report) + normalized.nodeid = _case_identity(report.nodeid) + return normalized + + +@pytest.hookimpl(wrapper=True, optionalhook=True) +def pytest_handlecrashitem(report: pytest.TestReport) -> Generator[None, None, None]: + result: Final = yield + report.nodeid = _case_identity(report.nodeid) # rebind-ok: xdist forwards this same report after the crash hook + return result + + def _owned(nodeid: str) -> bool: parts: Final = Path(nodeid.split("::", 1)[0]).parts return parts[:2] == ("tests", "integration") and len(parts) > 3 and parts[2] in OWNED_DIRECTORIES @@ -98,8 +159,13 @@ def _digest(seed: int, identity: str) -> bytes: return hashlib.sha256(f"{seed}:{identity}".encode()).digest() -def _order_key(seed: int, nodeid: str, group: str = "") -> tuple[bytes, bytes, bytes]: - return _digest(seed, nodeid.split("::", 1)[0]), _digest(seed, group or nodeid), _digest(seed, nodeid) +def _order_key(seed: int, nodeid: str, group: str = "", inner_group: str = "") -> tuple[bytes, bytes, bytes, bytes]: + return ( + _digest(seed, nodeid.split("::", 1)[0]), + _digest(seed, group or nodeid), + _digest(seed, inner_group or group or nodeid), + _digest(seed, nodeid), + ) def _shard_files(path: Path | None) -> frozenset[str] | None: @@ -116,11 +182,28 @@ def _shard_files(path: Path | None) -> frozenset[str] | None: @pytest.hookimpl(wrapper=True) def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> Iterator[None]: + if config.getoption("dist", default="no") == "loadgroup" or config.getoption("loadgroup", default=False): + automatic_groups: Final = scheduling_groups(tuple(collected_case(item) for item in items)) + for item in items: + fixture_group: Final = automatic_groups[item.nodeid] + if fixture_group and not tuple(item.iter_markers("xdist_group")): + item.add_marker( + pytest.mark.xdist_group("integration_fixture_" + hashlib.sha256(fixture_group.encode()).hexdigest()) + ) yield order_seed: Final = INTEGER_OPTION.validate_python(config.getoption("integration_order_seed")) if order_seed: - groups: Final = fixture_groups(tuple(collected_case(item) for item in items)) - items.sort(key=lambda item: _order_key(order_seed, item.nodeid, groups[item.nodeid])) + cases: Final = tuple(collected_case(item) for item in items) + groups: Final = scheduling_groups(cases) + inner_groups: Final = fixture_groups(cases) + items.sort( + key=lambda item: _order_key( + order_seed, + _case_identity(item.nodeid), + _case_identity(groups[item.nodeid]), + _case_identity(inner_groups[item.nodeid]), + ) + ) root: Final = Path(__file__).parent owned: Final = tuple( item @@ -134,7 +217,7 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item raise pytest.UsageError("Integration contracts are owned by CircleCI") for item in owned: item.add_marker(pytest.mark.integration) - config.stash[INVENTORY] = tuple(item.nodeid for item in owned) + config.stash[INVENTORY] = tuple(_case_identity(item.nodeid) for item in owned) count: Final = INTEGER_OPTION.validate_python(config.getoption("integration_shard_count")) index: Final = INTEGER_OPTION.validate_python(config.getoption("integration_shard_index")) if count < 1 or not 0 <= index < count: @@ -150,7 +233,7 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item deselected: Final = tuple(item for item in owned if item.nodeid.split("::", 1)[0] not in files) config.hook.pytest_deselected(items=deselected) items[:] = selected - config.stash[COLLECTED] = tuple(item.nodeid for item in items if item in owned) + config.stash[COLLECTED] = tuple(_case_identity(item.nodeid) for item in items if item in owned) def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: @@ -212,7 +295,7 @@ def gateway() -> Iterator[Gateway]: @pytest.fixture(scope="session") def shared_provider_server() -> Iterator[SharedProvider]: - if os.environ.get("PYTEST_XDIST_WORKER"): + if os.environ.get("PYTEST_XDIST_WORKER") and os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1": pytest.fail("the shared fake provider needs tests to run one at a time; this group runs under pytest-xdist") with shared_provider() as server: yield server diff --git a/tests/integration/database/test_harness_lifecycle.py b/tests/integration/database/test_harness_lifecycle.py index 91688eef263..bee5261ea0e 100644 --- a/tests/integration/database/test_harness_lifecycle.py +++ b/tests/integration/database/test_harness_lifecycle.py @@ -4,21 +4,23 @@ import asyncio import errno import os import re +import shutil import signal import socket import subprocess import sys import time from collections import Counter -from collections.abc import Callable +from collections.abc import Callable, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import AbstractContextManager, ExitStack from pathlib import Path from queue import SimpleQueue -from threading import Event from tempfile import TemporaryDirectory +from threading import Event from types import FrameType from typing import Final, Literal +from xml.etree import ElementTree import httpx import psutil @@ -40,6 +42,7 @@ from tests.integration._support.database_relay import ( ) from tests.integration._support.runtime import configure_executor from tests.integration._support.sharding import ShardExecution, shard_errors +from tests.integration._support.wire import Reply, wire_server class ExitRecord(TypedDict): @@ -53,6 +56,44 @@ RelayFactory = Callable[[str, bytes], AbstractContextManager[tuple[Relay, str]]] STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +def _nested_pytest_environment(environment: Mapping[str, str]) -> dict[str, str]: + inherited_run_settings: Final = frozenset( + ( + "INTEGRATION_RESULTS_DIR", + "INTEGRATION_ROUTING", + "INTEGRATION_WORKER_ISOLATION", + "INTEGRATION_CASE_SHARDS", + "INTEGRATION_CASE_SHARD_INDEX", + "CIRCLE_NODE_TOTAL", + "CIRCLE_NODE_INDEX", + "PYTEST_XDIST_WORKER", + "PYTEST_XDIST_WORKER_COUNT", + "PYTEST_XDIST_TESTRUNUID", + "GITHUB_ACTIONS", + ) + ) + return {name: value for name, value in environment.items() if name not in inherited_run_settings} + + +def test_owned_readiness_rejects_a_healthy_listener_from_another_process( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("INTEGRATION_PROXY_READY_SECONDS", "0") + with wire_server(lambda request: Reply(body=b"healthy")) as foreign: + with subprocess.Popen( + (sys.executable, "-I", "-c", "import time; time.sleep(30)"), start_new_session=True + ) as process: + try: + port: Final = int(foreign.url.rsplit(":", 1)[1]) + launch: Final = process_support._Launch(process, port, tmp_path / "proxy.log", time.monotonic()) + with pytest.raises(AssertionError, match="readiness deadline"): + process_support._wait_until_ready(launch) + assert process.poll() is None + finally: + process.kill() + process.wait(timeout=3) + + @pytest.mark.parametrize("factory", (database_relay, dropped_connection_relay, held_statement_relay)) def test_relay_closes_active_connections_and_releases_its_listener(factory: RelayFactory) -> None: with ExitStack() as cleanup: @@ -129,17 +170,29 @@ def test_slow_worker_startup_survives_healthchecks_and_a_stalled_worker_is_repla gateway: Gateway, tmp_path: Path ) -> None: hook: Final = tmp_path / "slow_worker_startup.py" - hook.write_text("from threading import Event\n\ndef configure():\n Event().wait(6)\n") + receipts: Final = tmp_path / "startup-receipts" + receipts.mkdir() + hook.write_text( + "import os\nfrom pathlib import Path\nfrom threading import Event\n\ndef configure():\n" + f" root = Path({str(receipts)!r})\n" + " try:\n" + " with (root / 'first').open('x'):\n" + " pass\n" + " except FileExistsError:\n" + " Event().wait(6)\n" + " (root / str(os.getpid())).touch()\n" + ) overrides: Final = { "PYTHONPATH": os.pathsep.join((str(tmp_path), os.environ.get("PYTHONPATH", ""))), "LITELLM_WORKER_STARTUP_HOOKS": "slow_worker_startup:configure", } with process_support.owned_proxy_process(gateway, tmp_path, overrides, workers=2) as owned: - workers: Final = eventually( - lambda: tuple(int(pid) for pid in STARTED_WORKER.findall(owned.log.read_text())), - lambda pids: len(pids) == 2 and owned.log.read_text().count("Application startup complete.") == 2, - seconds=30, + workers: Final = tuple(int(pid) for pid in STARTED_WORKER.findall(owned.log.read_text())) + assert len(workers) == 2, owned.log.read_text() + assert all((receipts / str(pid)).exists() for pid in workers), ( + "Proxy yielded before both workers finished startup" ) + assert owned.log.read_text().count("Application startup complete.") == 2 victim: Final = psutil.Process(workers[0]) victim.suspend() try: @@ -264,11 +317,7 @@ def test_seeded_collection_reuses_scoped_fixtures_and_executes_every_case(tmp_pa "--timeout=15", "-q", ), - env={ - name: value - for name, value in os.environ.items() - if name not in ("INTEGRATION_RESULTS_DIR", "INTEGRATION_ROUTING") - }, + env=_nested_pytest_environment(os.environ), capture_output=True, text=True, timeout=30, @@ -306,7 +355,6 @@ def test_native_file_selection_reuses_fixtures_and_qualifies_an_empty_shard( " yield 'shared'\n" " record('stop shared')\n" "@pytest.mark.parametrize('number', range(4))\n" - "@pytest.mark.xdist_group('shared')\n" "def test_shared(shared, number):\n" " assert shared == 'shared'\n" " record('shared ' + str(number))\n" @@ -358,22 +406,13 @@ def test_native_file_selection_reuses_fixtures_and_qualifies_an_empty_shard( "--integration-order-seed=1848224111", "--hypothesis-seed=4106601", "--timeout=15", + f"--junitxml={outputs[index] / 'junit.xml'}", *(("-n", str(workers), "--dist=loadgroup") if workers > 1 else ()), "-q", ), cwd=root, env={ - **{ - name: value - for name, value in os.environ.items() - if name - not in ( - "INTEGRATION_RESULTS_DIR", - "INTEGRATION_ROUTING", - "GITHUB_ACTIONS", - "PYTEST_XDIST_WORKER", - ) - }, + **_nested_pytest_environment(os.environ), "INTEGRATION_RESULTS_DIR": str(outputs[index]), "PYTHONPATH": os.pathsep.join((str(root), str(root / "tests"), str(root / "tests/e2e"))), }, @@ -391,6 +430,12 @@ def test_native_file_selection_reuses_fixtures_and_qualifies_an_empty_shard( ShardExecution.model_validate_json((output / "execution.json").read_text()) for output in outputs ) assert shard_errors(executions, 3) == () + for output, execution in zip(outputs, executions, strict=True): + cases: Final = ElementTree.parse(output / "junit.xml").findall(".//testcase") + assert sorted(case.attrib["name"] for case in cases) == sorted( + nodeid.rsplit("::", 1)[-1] for nodeid in execution.collected + ) + assert all("@integration_fixture_" not in nodeid for nodeid in execution.inventory) assert len(executions[0].inventory) == (2 if subset else 13) assert executions[2].collected == () assert executions[2].complete @@ -406,6 +451,235 @@ def test_native_file_selection_reuses_fixtures_and_qualifies_an_empty_shard( ) +@pytest.mark.parametrize(("workers", "seed"), ((1, 1848224111), (2, 0), (2, 984872314))) +@pytest.mark.parametrize("class_parameters", (False, True)) +def test_mixed_fixture_scopes_reuse_outer_and_inner_configurations( + tmp_path: Path, workers: int, seed: int, class_parameters: bool +) -> None: + events: Final = tmp_path / "events.txt" + root: Final = Path(__file__).resolve().parents[3] + with TemporaryDirectory(prefix="mixed-scopes-", dir=Path(__file__).parent) as directory: + probe: Final = Path(directory) / "test_fixture_probe.py" + probe.write_text( + "import pytest\nfrom pathlib import Path\n" + f"EVENTS = Path({str(events)!r})\n" + "def record(value):\n" + " with EVENTS.open('a') as output:\n" + " output.write(value + '\\n')\n" + "@pytest.fixture(scope='module', params=('M0', 'M1'))\n" + "def module_rig(request):\n" + " record('start ' + request.param)\n" + " yield request.param\n" + " record('stop ' + request.param)\n" + f"@pytest.fixture(scope='class', params={('C0', 'C1') if class_parameters else None!r})\n" + "def class_rig(request, module_rig):\n" + " identity = (\n" + " module_rig + '/' + request.cls.__name__\n" + " + ('/' + request.param if hasattr(request, 'param') else '')\n" + " )\n" + " record('start ' + identity)\n" + " yield identity\n" + " record('stop ' + identity)\n" + "class TestFirst:\n" + " @pytest.mark.parametrize('number', range(3))\n" + " def test_case(self, module_rig, class_rig, number):\n" + " assert class_rig.startswith(module_rig + '/')\n" + " record('case ' + class_rig + '/' + str(number))\n" + " def test_skip(self, class_rig):\n" + " pytest.skip('retained mixed-scope skip')\n" + "class TestSecond:\n" + " @pytest.mark.parametrize('number', range(3))\n" + " def test_case(self, module_rig, class_rig, number):\n" + " assert class_rig.startswith(module_rig + '/')\n" + " record('case ' + class_rig + '/' + str(number))\n" + ) + result: Final = subprocess.run( + ( + sys.executable, + "-P", + "tests/integration/run.py", + "database", + "--results", + str(tmp_path), + "--workers", + str(workers), + "--order-seed", + str(seed), + str(probe.relative_to(root)), + ), + cwd=root, + env=_nested_pytest_environment(os.environ), + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + execution: Final = ShardExecution.model_validate_json((tmp_path / "execution.json").read_text()) + assert execution.complete + expected_passes: Final = 24 if class_parameters else 12 + expected_skips: Final = 4 if class_parameters else 2 + assert len(execution.collected) == expected_passes + expected_skips + assert len(execution.passed) == expected_passes + assert len(execution.skipped) == expected_skips + cases: Final = ElementTree.parse(tmp_path / "junit.xml").findall(".//testcase") + assert sorted(case.attrib["name"] for case in cases) == sorted( + nodeid.rsplit("::", 1)[-1] for nodeid in execution.collected + ) + observed: Final = tuple(events.read_text().splitlines()) + starts: Final = Counter(event.removeprefix("start ") for event in observed if event.startswith("start ")) + stops: Final = Counter(event.removeprefix("stop ") for event in observed if event.startswith("stop ")) + assert starts == stops + assert len(starts) == (10 if class_parameters else 6) + assert all(count == 1 for count in starts.values()), observed + executed: Final = tuple(event for event in observed if event.startswith("case ")) + assert len(executed) == len(frozenset(executed)) == expected_passes + + +@pytest.mark.parametrize("workers", (1, 2)) +def test_direct_and_dependent_consumers_reuse_the_same_module_parameter(tmp_path: Path, workers: int) -> None: + source: Final = Path(__file__).resolve().parents[3] + root: Final = tmp_path / "probe-root" + for relative in ( + "pyproject.toml", + "tests/__init__.py", + "tests/integration/__init__.py", + "tests/integration/conftest.py", + "tests/integration/run.py", + ): + destination: Final = root / relative + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(source / relative, destination) + shutil.copytree( + source / "tests/integration/_support", + root / "tests/integration/_support", + ignore=shutil.ignore_patterns("__pycache__"), + ) + events: Final = tmp_path / "events.txt" + results: Final = tmp_path / "results" + probe: Final = root / "tests/integration/database/test_fixture_probe.py" + probe.parent.mkdir(parents=True, exist_ok=True) + probe.write_text( + "import pytest\nfrom pathlib import Path\n" + f"EVENTS = Path({str(events)!r})\n" + "def record(value):\n" + " with EVENTS.open('a') as output:\n" + " output.write(value + '\\n')\n" + "@pytest.fixture(scope='module', params=('A', 'B'))\n" + "def outer(request):\n" + " record('start ' + request.param)\n" + " yield request.param\n" + " record('stop ' + request.param)\n" + "@pytest.fixture(scope='module')\n" + "def dependent(outer):\n" + " return outer\n" + "@pytest.mark.parametrize('number', range(5))\n" + "def test_parent(outer, number):\n" + " record(f'case {outer}-parent-{number}')\n" + "@pytest.mark.parametrize('number', range(5))\n" + "def test_child(outer, dependent, number):\n" + " assert dependent == outer\n" + " record(f'case {outer}-child-{number}')\n" + "def test_retained_skip(outer):\n" + " pytest.skip('retained shared consumer skip')\n" + ) + result: Final = subprocess.run( + ( + sys.executable, + "-P", + str(root / "tests/integration/run.py"), + "database", + "--workers", + str(workers), + "--order-seed", + "1", + "--results", + str(results), + str(probe.relative_to(root)), + ), + cwd=root, + env=_nested_pytest_environment( + { + **os.environ, + "INTEGRATION_CASE_SHARDS": "2", + "INTEGRATION_CASE_SHARD_INDEX": "1", + "CIRCLE_NODE_TOTAL": "2", + "CIRCLE_NODE_INDEX": "1", + } + ), + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + execution: Final = ShardExecution.model_validate_json((results / "execution.json").read_bytes()) + assert execution.complete and len(execution.collected) == 22 + assert len(execution.passed) == 20 and len(execution.skipped) == 2 + assert all("::test_retained_skip[" in node for node in execution.skipped) + junit: Final = ElementTree.parse(results / "junit.xml") + assert len(junit.findall(".//testcase")) == 22 and len(junit.findall(".//skipped")) == 2 + assert all("@integration_fixture_" not in node.attrib["name"] for node in junit.findall(".//testcase")) + observed: Final = tuple(events.read_text().splitlines()) + assert Counter(event for event in observed if event.startswith("start ")) == Counter({"start A": 1, "start B": 1}) + assert Counter(event for event in observed if event.startswith("stop ")) == Counter({"stop A": 1, "stop B": 1}) + executed: Final = tuple(event for event in observed if event.startswith("case ")) + assert len(executed) == len(frozenset(executed)) == 20 + + +def test_worker_crash_keeps_the_original_junit_identity_and_an_incomplete_execution(tmp_path: Path) -> None: + root: Final = Path(__file__).resolve().parents[3] + with TemporaryDirectory(prefix="worker-crash-", dir=Path(__file__).parent) as directory: + probe: Final = Path(directory) / "test_fixture_probe.py" + probe.write_text( + "import os, pytest\n" + "@pytest.fixture(scope='module')\n" + "def shared():\n" + " yield 'shared'\n" + "def test_worker_crash(shared):\n" + " assert shared == 'shared'\n" + " os._exit(7)\n" + ) + result: Final = subprocess.run( + ( + sys.executable, + "-P", + "-m", + "pytest", + str(probe.relative_to(root)), + "--noconftest", + "-p", + "tests.integration.conftest", + "-o", + "addopts=", + "-n", + "2", + "--dist=loadgroup", + "--max-worker-restart=0", + f"--junitxml={tmp_path / 'junit.xml'}", + "-q", + ), + cwd=root, + env={ + **_nested_pytest_environment(os.environ), + "INTEGRATION_RESULTS_DIR": str(tmp_path), + }, + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert result.returncode == 1, result.stdout + result.stderr + execution: Final = ShardExecution.model_validate_json((tmp_path / "execution.json").read_text()) + assert not execution.complete + assert execution.passed == () + assert execution.skipped == () + assert execution.collected == (str(probe.relative_to(root)) + "::test_worker_crash",) + cases: Final = ElementTree.parse(tmp_path / "junit.xml").findall(".//testcase") + assert tuple(case.attrib.get("name") for case in cases) == ("test_worker_crash",) + assert cases[0].find("error") is not None + + @pytest.mark.parametrize("assignment", ("absent", "missing", "foreign")) def test_invalid_native_selections_fail_before_any_case_runs( tmp_path: Path, assignment: Literal["absent", "missing", "foreign"] @@ -437,11 +711,7 @@ def test_invalid_native_selections_fail_before_any_case_runs( ), cwd=root, env={ - **{ - name: value - for name, value in os.environ.items() - if name not in ("INTEGRATION_RESULTS_DIR", "INTEGRATION_ROUTING", "GITHUB_ACTIONS") - }, + **_nested_pytest_environment(os.environ), "PYTHONPATH": os.pathsep.join((str(root), str(root / "tests"), str(root / "tests/e2e"))), }, capture_output=True, @@ -453,9 +723,12 @@ def test_invalid_native_selections_fail_before_any_case_runs( assert not events.exists() -@pytest.mark.parametrize("entrypoint", ("gateway", "refused")) +@pytest.mark.parametrize("entrypoint", ("gateway", "refused", "zero-workers")) def test_owned_entrypoints_record_exit_and_reap_the_process_group( - gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, entrypoint: Literal["gateway", "refused"] + gateway: Gateway, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + entrypoint: Literal["gateway", "refused", "zero-workers"], ) -> None: monkeypatch.setenv("INTEGRATION_RESULTS_DIR", str(tmp_path)) config: Final = tmp_path / "config.yaml" @@ -471,7 +744,9 @@ def test_owned_entrypoints_record_exit_and_reap_the_process_group( "model_list: []\ngeneral_settings:\n master_key: os.environ/LITELLM_MASTER_KEY\n" " database_url: os.environ/DATABASE_URL\n" ) - with process_support.owned_gateway_image(gateway, tmp_path, {}, config=config, workers=1) as owned: + with process_support.owned_gateway_image( + gateway, tmp_path, {}, config=config, workers=0 if entrypoint == "zero-workers" else 1 + ) as owned: response: Final = owned.gateway.request("GET", "/health/readiness") assert response.status_code == 200, response.text completed: Final = TypeAdapter(ExitRecord).validate_json(owned.log.with_suffix(".exit.json").read_text()) diff --git a/tests/integration/management/test_auto_router_databricks_classifier_writes.py b/tests/integration/management/test_auto_router_databricks_classifier_writes.py index d954fd49d81..af486d90a7e 100644 --- a/tests/integration/management/test_auto_router_databricks_classifier_writes.py +++ b/tests/integration/management/test_auto_router_databricks_classifier_writes.py @@ -153,11 +153,11 @@ def _request_rows(rows: Sequence[dict[str, JsonValue]]) -> tuple[dict[str, JsonV return tuple(row for row in rows if row["origin"] != "autorouter_classifier") -def _rows(router: str, *, requests: int) -> tuple[dict[str, JsonValue], ...]: +def _rows(router: str, *, want: frozenset[str]) -> tuple[dict[str, JsonValue], ...]: return tuple( eventually( lambda: read_rows(_ROWS_QUERY, (router,)), - lambda found: len(_request_rows(found)) >= requests, + lambda found: want <= frozenset(string_value(row["request_id"]) for row in _request_rows(found)), seconds=70, ) ) @@ -395,7 +395,7 @@ async def test_moving_the_workspace_under_a_burst_moves_the_classifier_and_logs_ *tuple(convergence_calls.get_nowait() for _ in range(convergence_calls.qsize())), ) assert len(set(identities)) == len(identities), identities - rows: Final = _rows(name, requests=len(identities)) + rows: Final = _rows(name, want=frozenset(identities)) request_rows: Final = _request_rows(rows) assert sorted(string_value(row["request_id"]) for row in request_rows) == sorted(identities), request_rows assert all(row["status"] == "success" for row in rows), rows diff --git a/tests/integration/observability/conftest.py b/tests/integration/observability/conftest.py index fcac8eaba3c..101c2958b64 100644 --- a/tests/integration/observability/conftest.py +++ b/tests/integration/observability/conftest.py @@ -1,20 +1,21 @@ from __future__ import annotations -import uuid -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Iterator from pathlib import Path from typing import Final -from urllib.parse import urlparse import pytest -import yaml from integration._support.client import eventually -from integration._support.otlp_sink import SpanSinks, owned_sinks +from integration._support.otlp_sink import ( + AuditConfigWriter, + SpanSinks, + audit_configuration_writer, + langfuse_callback_variables, + owned_sinks, +) from integration._support.prometheus_series import CapRig, series_cap_rig, spend_rows from pydantic import JsonValue -AuditConfigWriter = Callable[[Path, Mapping[str, JsonValue]], Path] - @pytest.fixture(scope="module") def audit_sinks(tmp_path_factory: pytest.TempPathFactory) -> Iterator[SpanSinks]: @@ -25,34 +26,12 @@ def audit_sinks(tmp_path_factory: pytest.TempPathFactory) -> Iterator[SpanSinks] @pytest.fixture(scope="module") def otel_audit_config(audit_sinks: SpanSinks) -> AuditConfigWriter: - tenant_host: Final = urlparse(audit_sinks.tenant).netloc - - def write(directory: Path, litellm_settings: Mapping[str, JsonValue] = {}) -> Path: - config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) - config["litellm_settings"] = { - **config.get("litellm_settings", {}), - "callbacks": ["otel"], - "provider_url_destination_allowed_hosts": [tenant_host], - **dict(litellm_settings), - } - config["callback_settings"] = { - "otel": {"exporter": "http/json", "endpoint": audit_sinks.operator, "use_simple_processor": True} - } - config["general_settings"] = {**config.get("general_settings", {}), "user_api_key_cache_ttl": 2} - path: Final = directory / f"otel-audit-{uuid.uuid4().hex}.yaml" - path.write_text(yaml.safe_dump(config)) - return path - - return write + return audit_configuration_writer(audit_sinks) @pytest.fixture(scope="module") def langfuse_vars(audit_sinks: SpanSinks) -> dict[str, JsonValue]: - return { - "langfuse_public_key": "pk-lf-audit", - "langfuse_secret_key": "sk-lf-audit", - "langfuse_host": audit_sinks.tenant, - } + return langfuse_callback_variables(audit_sinks) @pytest.fixture(scope="session") diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_modes.py b/tests/integration/observability/test_arize_otel_v2_openinference_modes.py index 7d0f3f31787..cb61ccf2ac1 100644 --- a/tests/integration/observability/test_arize_otel_v2_openinference_modes.py +++ b/tests/integration/observability/test_arize_otel_v2_openinference_modes.py @@ -411,6 +411,7 @@ def test_arize_otel_v2_c9_upstream_failures_are_recorded(failure_status: int, ga assert not any(".tool_calls." in key for key in attributes), attributes +@pytest.mark.timeout(180) def test_arize_otel_v2_c10_attribute_limit_keeps_tool_prefix(gateway: Gateway, tmp_path: Path) -> None: marker: Final = "c10-" + uuid.uuid4().hex calls: Final = _response_tool_calls(marker, ("Paris", "Berlin", "Rome", "Tokyo", "Oslo", "Lima", "Accra", "Delhi")) diff --git a/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py b/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py index dc7c8beba66..5d7b5948d96 100644 --- a/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py +++ b/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py @@ -422,6 +422,7 @@ def test_x3_killing_one_worker_leaves_the_other_serving( _assert_spend(response.headers["x-litellm-call-id"], response.text) +@pytest.mark.timeout(180) def test_x4_proxy_restart_preserves_completed_spend_rows( gateway: Gateway, tmp_path: Path ) -> None: diff --git a/tests/integration/observability/test_azure_storage_chaos.py b/tests/integration/observability/test_azure_storage_chaos.py index 079ce72f9ba..2235dc8fc56 100644 --- a/tests/integration/observability/test_azure_storage_chaos.py +++ b/tests/integration/observability/test_azure_storage_chaos.py @@ -4,6 +4,8 @@ import uuid from pathlib import Path from typing import Final +import pytest + import httpx from _azure_storage_support import ( SINK_HOSTS, @@ -174,6 +176,7 @@ def test_killing_one_worker_keeps_the_other_serving_and_uploading(gateway: Gatew ) +@pytest.mark.timeout(180) def test_restarting_the_proxy_before_the_queue_flushes_bounds_the_loss_to_the_unflushed_queue_and_recovers( gateway: Gateway, tmp_path: Path ) -> None: diff --git a/tests/integration/observability/test_bedrock_guardrail_attachments.py b/tests/integration/observability/test_bedrock_guardrail_attachments.py index 0d53e335cd1..aba68544216 100644 --- a/tests/integration/observability/test_bedrock_guardrail_attachments.py +++ b/tests/integration/observability/test_bedrock_guardrail_attachments.py @@ -6,7 +6,7 @@ from collections.abc import Callable, Iterator, Mapping from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path -from typing import Final +from typing import Final, Literal import httpx import pytest @@ -39,9 +39,7 @@ def _b64(raw: bytes) -> str: @dataclass(frozen=True, slots=True) class _Rig: - strict: Gateway - lenient: Gateway - checks: Gateway + gateway: Gateway policy: Wire upstream: Wire outage: threading.Event @@ -212,34 +210,60 @@ def _config( return path +@dataclass(frozen=True, slots=True) +class _Wires: + directory: Path + policy: Wire + upstream: Wire + outage: threading.Event + + @pytest.fixture(scope="module") -def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: +def wires(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Wires]: directory: Final = tmp_path_factory.mktemp("bedrock-guardrail-attachments") outage: Final = threading.Event() + with wire_server(_guardrail(outage)) as policy, wire_server(_upstream) as upstream: + yield _Wires(directory, policy, upstream, outage) + + +def _attachment_proxy(wires: _Wires, mode: Literal["strict", "lenient", "checks"]) -> Iterator[_Rig]: with ( - wire_server(_guardrail(outage)) as policy, - wire_server(_upstream) as upstream, gateway_from_environment() as shared, owned_proxy( - shared, directory, {}, config=_config(directory, "strict", policy, upstream, False), workers=2 - ) as strict, - owned_proxy( - shared, directory, {}, config=_config(directory, "lenient", policy, upstream, True), workers=2 - ) as lenient, - owned_proxy( - shared, directory, {}, config=_config(directory, "checks", policy, upstream, False, checks=True), workers=2 - ) as checks, + shared, + wires.directory, + {}, + config=_config( + wires.directory, mode, wires.policy, wires.upstream, mode == "lenient", checks=mode == "checks" + ), + workers=2, + ) as gateway, ): - yield _Rig(strict, lenient, checks, policy, upstream, outage) + yield _Rig(gateway, wires.policy, wires.upstream, wires.outage) + + +@pytest.fixture(scope="module") +def rig(wires: _Wires) -> Iterator[_Rig]: + yield from _attachment_proxy(wires, "strict") + + +@pytest.fixture(scope="module") +def lenient_rig(wires: _Wires) -> Iterator[_Rig]: + yield from _attachment_proxy(wires, "lenient") + + +@pytest.fixture(scope="module") +def checks_rig(wires: _Wires) -> Iterator[_Rig]: + yield from _attachment_proxy(wires, "checks") @pytest.fixture(autouse=True) -def _drained(rig: _Rig) -> Iterator[None]: - rig.policy.drain() - rig.upstream.drain() - rig.outage.clear() +def _drained(wires: _Wires) -> Iterator[None]: + wires.policy.drain() + wires.upstream.drain() + wires.outage.clear() yield - rig.outage.clear() + wires.outage.clear() def _image_block(endpoint: str, encoded: str, mime: str = "image/png") -> dict[str, JsonValue]: @@ -306,7 +330,7 @@ def _forwarded(upstream: Wire) -> list[dict[str, JsonValue]]: @pytest.mark.parametrize("endpoint", _ENDPOINTS) def test_an_image_only_turn_is_scanned_and_a_threat_image_is_blocked(rig: _Rig, endpoint: str) -> None: - response: Final = _send(rig.strict, endpoint, _image_block(endpoint, _b64(_THREAT_PNG))) + response: Final = _send(rig.gateway, endpoint, _image_block(endpoint, _b64(_THREAT_PNG))) assert response.status_code == 400, response.text assert "VIOLENCE" in response.text, response.text @@ -316,7 +340,7 @@ def test_an_image_only_turn_is_scanned_and_a_threat_image_is_blocked(rig: _Rig, @pytest.mark.parametrize("endpoint", _ENDPOINTS) def test_a_harmless_image_is_scanned_and_reaches_the_model(rig: _Rig, endpoint: str) -> None: - response: Final = _send(rig.strict, endpoint, _image_block(endpoint, _b64(_PLAIN_PNG))) + response: Final = _send(rig.gateway, endpoint, _image_block(endpoint, _b64(_PLAIN_PNG))) assert response.status_code == 200, response.text assert _image_scans(_scans(rig.policy)) == [[{"format": "png", "source": {"bytes": _b64(_PLAIN_PNG)}}]] @@ -326,7 +350,7 @@ def test_a_harmless_image_is_scanned_and_reaches_the_model(rig: _Rig, endpoint: @pytest.mark.parametrize("endpoint", _ENDPOINTS) def test_a_pdf_is_refused_without_calling_bedrock_or_the_model(rig: _Rig, endpoint: str) -> None: question: Final = f"What is in the attachment {uuid.uuid4().hex}?" - response: Final = _send(rig.strict, endpoint, _text_block(endpoint, question), _pdf_block(endpoint)) + response: Final = _send(rig.gateway, endpoint, _text_block(endpoint, question), _pdf_block(endpoint)) assert response.status_code == 400, response.text assert "Bedrock guardrail cannot scan 1 attachment(s)" in response.text, response.text @@ -335,39 +359,41 @@ def test_a_pdf_is_refused_without_calling_bedrock_or_the_model(rig: _Rig, endpoi @pytest.mark.parametrize("endpoint", _ENDPOINTS) -def test_skip_unscannable_attachments_lets_a_pdf_through_and_still_scans_the_text(rig: _Rig, endpoint: str) -> None: +def test_skip_unscannable_attachments_lets_a_pdf_through_and_still_scans_the_text( + lenient_rig: _Rig, endpoint: str +) -> None: question: Final = f"What is in the attachment {uuid.uuid4().hex}?" - response: Final = _send(rig.lenient, endpoint, _text_block(endpoint, question), _pdf_block(endpoint)) + response: Final = _send(lenient_rig.gateway, endpoint, _text_block(endpoint, question), _pdf_block(endpoint)) assert response.status_code == 200, response.text - assert question in _text_scans(_scans(rig.policy)) - (forwarded,) = _forwarded(rig.upstream) + assert question in _text_scans(_scans(lenient_rig.policy)) + (forwarded,) = _forwarded(lenient_rig.upstream) assert _b64(_PDF) in json.dumps(forwarded) @pytest.mark.parametrize("endpoint", _ENDPOINTS) -def test_skip_unscannable_attachments_still_blocks_a_threat_image(rig: _Rig, endpoint: str) -> None: - response: Final = _send(rig.lenient, endpoint, _image_block(endpoint, _b64(_THREAT_PNG))) +def test_skip_unscannable_attachments_still_blocks_a_threat_image(lenient_rig: _Rig, endpoint: str) -> None: + response: Final = _send(lenient_rig.gateway, endpoint, _image_block(endpoint, _b64(_THREAT_PNG))) assert response.status_code == 400, response.text assert "VIOLENCE" in response.text, response.text - assert _forwarded(rig.upstream) == [] + assert _forwarded(lenient_rig.upstream) == [] -def test_skip_unscannable_attachments_still_scans_a_converse_pdf_name_and_context(rig: _Rig) -> None: +def test_skip_unscannable_attachments_still_scans_a_converse_pdf_name_and_context(lenient_rig: _Rig) -> None: context: Final = f"Employee SSN {_SSN} {uuid.uuid4().hex}" document: Final = { "document": {"format": "pdf", "name": "record", "context": context, "source": {"bytes": _b64(_PDF)}} } - response: Final = _send(rig.lenient, "converse", _text_block("converse", "summarize"), document) + response: Final = _send(lenient_rig.gateway, "converse", _text_block("converse", "summarize"), document) assert response.status_code == 400, response.text - assert f"record\n{context}" in _text_scans(_scans(rig.policy)) - assert _forwarded(rig.upstream) == [] + assert f"record\n{context}" in _text_scans(_scans(lenient_rig.policy)) + assert _forwarded(lenient_rig.upstream) == [] def test_a_streaming_chat_request_with_a_pdf_is_refused_before_streaming(rig: _Rig) -> None: - response: Final = _send(rig.strict, "chat", _text_block("chat", "summarize"), _pdf_block("chat"), stream=True) + response: Final = _send(rig.gateway, "chat", _text_block("chat", "summarize"), _pdf_block("chat"), stream=True) assert response.status_code == 400, response.text assert "Bedrock guardrail cannot scan 1 attachment(s)" in response.text, response.text @@ -398,7 +424,7 @@ def test_a_text_document_with_an_ssn_is_scanned_as_text_and_blocked( rig: _Rig, endpoint: str, document: Callable[[str], dict[str, JsonValue]] ) -> None: secret: Final = f"Employee SSN {_SSN} {uuid.uuid4().hex}" - response: Final = _send(rig.strict, endpoint, _text_block(endpoint, "summarize"), document(secret)) + response: Final = _send(rig.gateway, endpoint, _text_block(endpoint, "summarize"), document(secret)) assert response.status_code == 400, response.text assert any(secret in text for text in _text_scans(_scans(rig.policy))) @@ -410,7 +436,7 @@ def test_a_harmless_text_document_is_scanned_as_text_and_reaches_the_model( rig: _Rig, endpoint: str, document: Callable[[str], dict[str, JsonValue]] ) -> None: note: Final = f"The meeting is at noon {uuid.uuid4().hex}" - response: Final = _send(rig.strict, endpoint, _text_block(endpoint, "summarize"), document(note)) + response: Final = _send(rig.gateway, endpoint, _text_block(endpoint, "summarize"), document(note)) assert response.status_code == 200, response.text assert any(note in text for text in _text_scans(_scans(rig.policy))) @@ -420,7 +446,7 @@ def test_a_harmless_text_document_is_scanned_as_text_and_reaches_the_model( def test_a_converse_txt_document_sent_as_bytes_stays_unscannable(rig: _Rig) -> None: document: Final = {"document": {"format": "txt", "name": "record", "source": {"bytes": _b64(b"notes")}}} - response: Final = _send(rig.strict, "converse", _text_block("converse", "summarize"), document) + response: Final = _send(rig.gateway, "converse", _text_block("converse", "summarize"), document) assert response.status_code == 400, response.text assert "cannot scan 1 attachment(s) (document)" in response.text, response.text @@ -429,7 +455,7 @@ def test_a_converse_txt_document_sent_as_bytes_stays_unscannable(rig: _Rig) -> N def test_a_converse_guard_content_threat_image_is_scanned_and_blocked(rig: _Rig) -> None: guarded: Final = {"guardContent": {"image": {"format": "png", "source": {"bytes": _b64(_THREAT_PNG)}}}} - response: Final = _send(rig.strict, "converse", _text_block("converse", "describe"), guarded) + response: Final = _send(rig.gateway, "converse", _text_block("converse", "describe"), guarded) assert response.status_code == 400, response.text assert "VIOLENCE" in response.text, response.text @@ -437,24 +463,24 @@ def test_a_converse_guard_content_threat_image_is_scanned_and_blocked(rig: _Rig) @pytest.mark.parametrize("endpoint", _ENDPOINTS) -def test_checks_mode_refuses_an_image_without_calling_bedrock(rig: _Rig, endpoint: str) -> None: - response: Final = _send(rig.checks, endpoint, _image_block(endpoint, _b64(_PLAIN_PNG))) +def test_checks_mode_refuses_an_image_without_calling_bedrock(checks_rig: _Rig, endpoint: str) -> None: + response: Final = _send(checks_rig.gateway, endpoint, _image_block(endpoint, _b64(_PLAIN_PNG))) assert response.status_code == 400, response.text assert "cannot scan 1 attachment(s) (image)" in response.text, response.text - assert list(rig.policy.drain()) == [] - assert _forwarded(rig.upstream) == [] + assert list(checks_rig.policy.drain()) == [] + assert _forwarded(checks_rig.upstream) == [] -def test_checks_mode_sends_a_converse_text_document_to_invoke_guardrail_checks(rig: _Rig) -> None: +def test_checks_mode_sends_a_converse_text_document_to_invoke_guardrail_checks(checks_rig: _Rig) -> None: note: Final = f"The meeting is at noon {uuid.uuid4().hex}" - response: Final = _send(rig.checks, "converse", _converse_text_document({"text": note})) + response: Final = _send(checks_rig.gateway, "converse", _converse_text_document({"text": note})) assert response.status_code == 200, response.text - calls: Final = list(rig.policy.drain()) + calls: Final = list(checks_rig.policy.drain()) assert [call.target for call in calls] == [_CHECKS_PATH] assert note in calls[0].body.decode() - (forwarded,) = _forwarded(rig.upstream) + (forwarded,) = _forwarded(checks_rig.upstream) assert note in json.dumps(forwarded) @@ -466,14 +492,14 @@ def test_checks_mode_sends_a_converse_text_document_to_invoke_guardrail_checks(r ], ) def test_non_standard_base64_is_normalized_before_the_scan(rig: _Rig, encoded: str) -> None: - response: Final = _send(rig.strict, "chat", _image_block("chat", encoded)) + response: Final = _send(rig.gateway, "chat", _image_block("chat", encoded)) assert response.status_code == 400, response.text assert _image_scans(_scans(rig.policy)) == [[{"format": "png", "source": {"bytes": _b64(_THREAT_PNG)}}]] def test_the_image_format_sent_to_bedrock_comes_from_the_bytes(rig: _Rig) -> None: - response: Final = _send(rig.strict, "chat", _image_block("chat", _b64(_PLAIN_PNG), mime="image/jpeg")) + response: Final = _send(rig.gateway, "chat", _image_block("chat", _b64(_PLAIN_PNG), mime="image/jpeg")) assert response.status_code == 200, response.text assert _image_scans(_scans(rig.policy)) == [[{"format": "png", "source": {"bytes": _b64(_PLAIN_PNG)}}]] @@ -482,7 +508,7 @@ def test_the_image_format_sent_to_bedrock_comes_from_the_bytes(rig: _Rig) -> Non @pytest.mark.parametrize("endpoint", ("chat", "converse")) def test_a_guardrail_outage_fails_closed_on_an_image_turn(rig: _Rig, endpoint: str) -> None: rig.outage.set() - response: Final = _send(rig.strict, endpoint, _image_block(endpoint, _b64(_PLAIN_PNG))) + response: Final = _send(rig.gateway, endpoint, _image_block(endpoint, _b64(_PLAIN_PNG))) assert response.status_code == 500, response.text assert "synthetic guardrail outage" in response.text, response.text @@ -500,17 +526,17 @@ def test_a_concurrent_mixed_burst_across_two_workers_keeps_each_verdict(rig: _Ri def send(case: tuple[str, str]) -> tuple[str, str, int]: kind, tag = case if kind == "threat": - response = _send(rig.strict, "chat", _text_block("chat", tag), _image_block("chat", _b64(_THREAT_PNG))) + response = _send(rig.gateway, "chat", _text_block("chat", tag), _image_block("chat", _b64(_THREAT_PNG))) elif kind == "plain": response = _send( - rig.strict, "messages", _text_block("messages", tag), _image_block("messages", _b64(_PLAIN_PNG)) + rig.gateway, "messages", _text_block("messages", tag), _image_block("messages", _b64(_PLAIN_PNG)) ) elif kind == "pdf": - response = _send(rig.strict, "responses", _text_block("responses", tag), _pdf_block("responses")) + response = _send(rig.gateway, "responses", _text_block("responses", tag), _pdf_block("responses")) elif kind == "converse-note": - response = _send(rig.strict, "converse", _converse_text_document({"text": f"note {tag}"})) + response = _send(rig.gateway, "converse", _converse_text_document({"text": f"note {tag}"})) else: - response = _send(rig.strict, "converse", _converse_text_document({"text": f"SSN {_SSN} {tag}"})) + response = _send(rig.gateway, "converse", _converse_text_document({"text": f"SSN {_SSN} {tag}"})) return kind, tag, response.status_code with ThreadPoolExecutor(max_workers=12) as pool: diff --git a/tests/integration/observability/test_guardrail_stream_scope_matrix.py b/tests/integration/observability/test_guardrail_stream_scope_matrix.py index b5c66b202cc..8ab724d35fb 100644 --- a/tests/integration/observability/test_guardrail_stream_scope_matrix.py +++ b/tests/integration/observability/test_guardrail_stream_scope_matrix.py @@ -502,6 +502,7 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[MatrixRig]: "ANTHROPIC_API_KEY": "synthetic-anthropic-key", "GEMINI_API_BASE": provider.url, "GEMINI_API_KEY": "synthetic-gemini-key", + "LITELLM_BUILTIN_PASS_THROUGH_ROUTES_FIRST": "true", }, "general_settings": { **config_base["general_settings"], @@ -1575,6 +1576,7 @@ def test_eh_provider_passthrough_routes_classify_effective_streaming( "POST", route, request_body, + key=key, headers=headers, ) assert response.status_code == 200 and marker in response.text, response.text diff --git a/tests/integration/observability/test_langfuse_delivery.py b/tests/integration/observability/test_langfuse_delivery.py index 00eee082c33..018c6a35f31 100644 --- a/tests/integration/observability/test_langfuse_delivery.py +++ b/tests/integration/observability/test_langfuse_delivery.py @@ -9,6 +9,7 @@ from collections.abc import Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path +from queue import SimpleQueue from typing import Final import anthropic @@ -1723,12 +1724,17 @@ def test_audit_sink_outage_mid_burst_loses_no_spend_rows(gateway: Gateway, tmp_p baggage_session: Final = "baggage-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit outage: Final = threading.Event() + rejected_export: Final = threading.Event() + accepted: Final[SimpleQueue[Request]] = SimpleQueue() def langfuse(request: Request) -> Reply: if outage.is_set(): + if request.method == "POST" and request.target == TRACES_PATH: + rejected_export.set() return Reply(status=503) if request.method == "GET" and request.target.startswith(PROJECTS_PATH): return _projects() + accepted.put(request) return Reply(body=b"", content_type="application/x-protobuf") with ( @@ -1767,6 +1773,7 @@ def test_audit_sink_outage_mid_burst_loses_no_spend_rows(gateway: Gateway, tmp_p outage.set() with ThreadPoolExecutor(max_workers=30) as pool: first_wave: Final = tuple(pool.map(fire, jobs[:10])) + assert rejected_export.wait(timeout=10), "the exporter did not reach the scripted outage" unhealthy: Final = candidate.request("GET", "/health/services?service=langfuse") assert unhealthy.status_code != 200 or "unhealthy" in unhealthy.text, ( f"langfuse must report unhealthy while the sink 503s: {unhealthy.status_code} {unhealthy.text}" @@ -1786,7 +1793,7 @@ def test_audit_sink_outage_mid_burst_loses_no_spend_rows(gateway: Gateway, tmp_p received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones delivered: Final = eventually( lambda: ( - received.extend(destination.drain()) or tuple( + received.extend(accepted.get_nowait() for _ in range(accepted.qsize())) or tuple( span for span in _spans(received) if _attribute(span.attributes, "langfuse.observation.metadata.litellm_call_id") diff --git a/tests/integration/observability/test_logging_only_scope_chaos.py b/tests/integration/observability/test_logging_only_scope_chaos.py index d70073a5071..a34b13f048f 100644 --- a/tests/integration/observability/test_logging_only_scope_chaos.py +++ b/tests/integration/observability/test_logging_only_scope_chaos.py @@ -50,6 +50,7 @@ from openai import APIConnectionError as OpenAIAPIConnectionError from pydantic import JsonValue +@pytest.mark.timeout(180) def test_K1_policy_edge_restart_mid_burst_keeps_output_observation_fail_open(gateway: Gateway, tmp_path: Path) -> None: identity: Final = f"logging-scope-k1-{uuid.uuid4().hex}" marker: Final = uuid.uuid4().hex @@ -209,6 +210,7 @@ def test_K1_policy_edge_restart_mid_burst_keeps_output_observation_fail_open(gat ) == tuple((identity, "logging_only", expected_status) for _ in expected_directions), (index, entries) +@pytest.mark.timeout(180) def test_K2_policy_edge_delay_does_not_delay_concurrent_callers(gateway: Gateway, tmp_path: Path) -> None: identity: Final = f"logging-scope-k2-{uuid.uuid4().hex}" marker: Final = uuid.uuid4().hex diff --git a/tests/integration/observability/test_otel_excluded_services.py b/tests/integration/observability/test_otel_excluded_services.py index 09f40630391..1de071e7103 100644 --- a/tests/integration/observability/test_otel_excluded_services.py +++ b/tests/integration/observability/test_otel_excluded_services.py @@ -16,20 +16,38 @@ from integration._support.client import ( gateway_from_environment, ) from integration._support.otlp_sink import ( + AuditConfigWriter, Span, SpanSinks, + audit_configuration_writer, + langfuse_callback_variables, + owned_sinks, recorded_spans, spans_for_trace, ) from integration._support.process import owned_proxy, owned_proxy_process from pydantic import JsonValue -AuditConfigWriter = Callable[[Path, Mapping[str, JsonValue]], Path] - DB_SYSTEM_KEYS: Final = frozenset({"db.system.name", "db.system"}) -@pytest.fixture(scope="module") +@pytest.fixture +def audit_sinks(tmp_path: Path) -> Iterator[SpanSinks]: + with owned_sinks(tmp_path / "sinks") as sinks: + yield sinks + + +@pytest.fixture +def otel_audit_config(audit_sinks: SpanSinks) -> AuditConfigWriter: + return audit_configuration_writer(audit_sinks) + + +@pytest.fixture +def langfuse_vars(audit_sinks: SpanSinks) -> dict[str, JsonValue]: + return langfuse_callback_variables(audit_sinks) + + +@pytest.fixture def gateway(audit_sinks: SpanSinks) -> Iterator[Gateway]: with gateway_from_environment() as base: yield base diff --git a/tests/integration/observability/test_passthrough_upstream_error_chaos.py b/tests/integration/observability/test_passthrough_upstream_error_chaos.py index 568a0990c56..014942ccfaf 100644 --- a/tests/integration/observability/test_passthrough_upstream_error_chaos.py +++ b/tests/integration/observability/test_passthrough_upstream_error_chaos.py @@ -3,7 +3,8 @@ import json import re import signal import threading -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Final @@ -59,6 +60,16 @@ def _held_connections(pid: int, port: int) -> int: ) +@contextmanager +def _suspended_worker(pid: int) -> Iterator[None]: + worker: Final = psutil.Process(pid) + worker.suspend() + try: + yield + finally: + worker.resume() + + def _error_information(call_id: str) -> dict[str, JsonValue]: rows: Final = eventually( lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)), @@ -90,7 +101,7 @@ def _single_spend_row(call_id: str) -> None: async def _fire_burst( - base_url: str, key: str, count: int, *, tolerate_transport_errors: bool = False + base_url: str, key: str, count: int, *, start_index: int = 0, tolerate_transport_errors: bool = False ) -> tuple[_Served, ...]: async def one(client: httpx.AsyncClient, index: int) -> _Served: if index % 3 == 0: @@ -111,7 +122,7 @@ async def _fire_burst( async with httpx.AsyncClient(base_url=base_url, timeout=30, trust_env=False) as client: results: Final = await asyncio.gather( - *(one(client, index) for index in range(count)), return_exceptions=tolerate_transport_errors + *(one(client, index) for index in range(start_index, start_index + count)), return_exceptions=tolerate_transport_errors ) for result in results: assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) @@ -163,11 +174,23 @@ async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gat seconds=30, ) async with asyncio.TaskGroup() as tasks: - burst: Final = tasks.create_task( - _fire_burst(str(candidate.client.base_url), candidate.key, 20, tolerate_transport_errors=True) - ) try: - await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 20, 30) + with _suspended_worker(workers[0]): + seed: Final = tasks.create_task( + _fire_burst(str(candidate.client.base_url), candidate.key, 1, tolerate_transport_errors=True) + ) + await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 1, 30) + with _suspended_worker(workers[1]): + rest: Final = tasks.create_task( + _fire_burst( + str(candidate.client.base_url), + candidate.key, + 19, + start_index=1, + tolerate_transport_errors=True, + ) + ) + await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 20, 30) peer_port: Final = int(wire.url.rsplit(":", 1)[1]) held_by: Final = {pid: _held_connections(pid, peer_port) for pid in workers} assert sum(held_by.values()) == 20 and min(held_by.values()) > 0, held_by @@ -177,9 +200,10 @@ async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gat connection.raddr.port for connection in victim.net_connections(kind="tcp") if connection.raddr ) victim.send_signal(signal.SIGKILL) + victim.wait(timeout=5) finally: release.set() - served: Final = await burst + served: Final = (*await seed, *await rest) for item in served: assert item.response.status_code in (200, 404, 500, 502), item.response.status_code follow_up: Final = candidate.request( diff --git a/tests/integration/observability/test_s3_v2_partition_granularity.py b/tests/integration/observability/test_s3_v2_partition_granularity.py index dc0a7184cf8..a362ff55245 100644 --- a/tests/integration/observability/test_s3_v2_partition_granularity.py +++ b/tests/integration/observability/test_s3_v2_partition_granularity.py @@ -1166,6 +1166,7 @@ def test_s3_v2_hour_worker_kill_mid_burst_keeps_the_other_worker_logging(gateway assert _outside_layout(objects, "hour") == () +@pytest.mark.timeout(180) def test_s3_v2_hour_proxy_restart_mid_burst_keeps_the_layout_without_duplicates( gateway: Gateway, tmp_path: Path ) -> None: diff --git a/tests/integration/providers/test_device_code_login_guard_boot.py b/tests/integration/providers/test_device_code_login_guard_boot.py index ecec6402aee..687a8f1d313 100644 --- a/tests/integration/providers/test_device_code_login_guard_boot.py +++ b/tests/integration/providers/test_device_code_login_guard_boot.py @@ -241,11 +241,14 @@ def test_booting_without_tokens_drops_the_login_deployments_and_a_reload_after_m workers=2, ) as owned: workers: Final = _workers(owned) + _every_worker_lists(owned, workers, login_models=False) boot_log: Final = owned.log.read_text() assert peers.auth_connections() == () assert dl.CHATGPT_USER_CODE not in boot_log and dl.COPILOT_USER_CODE not in boot_log, boot_log[-3000:] - assert sorted(_DROPPED.findall(boot_log)) == sorted([*_FIXED.values(), *_WILDCARDS] * 2), boot_log[-3000:] - _every_worker_lists(owned, workers, login_models=False) + dropped_models: Final = _DROPPED.findall(boot_log) + expected_drops: Final = frozenset((*_FIXED.values(), *_WILDCARDS)) + assert frozenset(dropped_models) == expected_drops, boot_log[-3000:] + assert all(dropped_models.count(model) >= len(workers) for model in expected_drops), boot_log[-3000:] with owned.gateway.scenario() as scenario: key: Final = scenario.key() for name in _LOGIN_NAMES: diff --git a/tests/integration/providers/test_strands_decider_agentcore_wire.py b/tests/integration/providers/test_strands_decider_agentcore_wire.py index fc62849db9b..c7213a66bc5 100644 --- a/tests/integration/providers/test_strands_decider_agentcore_wire.py +++ b/tests/integration/providers/test_strands_decider_agentcore_wire.py @@ -418,7 +418,7 @@ def test_an_arn_deployment_answers_on_every_decisions_route_signed_for_bedrock_a assert call.headers[_SESSION_HEADER] == _default_session("audit_signed"), dict(call.headers) _assert_signed(call) for response, route in zip(responses, _ROUTES, strict=True): - row = _spend_row(response.headers["x-litellm-call-id"]) + row: Final = _spend_row(response.headers["x-litellm-call-id"]) assert row == { "status": "success", "call_type": _ROUTE_CALL_TYPES[route.path], diff --git a/tests/integration/run.py b/tests/integration/run.py index 98b727efce9..7db0087a94c 100644 --- a/tests/integration/run.py +++ b/tests/integration/run.py @@ -88,6 +88,11 @@ def main() -> int: "LITELLM_LOCAL_MODEL_COST_MAP": "True", "LITELLM_MODE": "PRODUCTION", "PYTHON_DOTENV_DISABLED": "1", + **( + {"INTEGRATION_WORKER_ISOLATION": "1"} + if options.workers > 1 and options.group in {"management", "accounting", "providers", "extensions"} + else {} + ), } result: Final = subprocess.call( [ @@ -114,8 +119,7 @@ def main() -> int: "-o", "junit_family=xunit1", *(("-n", str(options.workers)) if options.workers > 1 else ()), - *(("--dist=worksteal",) if options.group == "security" and options.workers > 1 else ()), - *(("--dist=loadgroup",) if options.group == "mcp" and options.workers > 1 else ()), + *(("--dist=loadgroup",) if options.workers > 1 else ()), ], cwd=root, env=environment, diff --git a/tests/integration/security/test_sweep_sensitivity.py b/tests/integration/security/test_sweep_sensitivity.py index 5c38069a796..46c004ad45d 100644 --- a/tests/integration/security/test_sweep_sensitivity.py +++ b/tests/integration/security/test_sweep_sensitivity.py @@ -19,15 +19,10 @@ from typing import Final import httpx import pytest -from redis import Redis from integration._support.client import Gateway, eventually, string_value +from integration._support.wire import Reply, Request, wire_server from integration.security._canary import DECODE_BUDGET_BYTES, MARKER, SLOTS, DecodeBudgetExceeded, canary, find_canary from integration.security._sinks import CONFIG_MODEL, GENERIC_SINK, Rig, canary_rig, chat_upstream, settle, team_caller -from integration._support.wire import Reply, Request, wire_server -from tests.integration._support.database import read_rows, write_rows -from tests.integration._support.redis_process import owned_redis -from tests.integration._support.tls import server_context, write_self_signed_cert -from tests.integration._support.workers import worker_services from integration.security._sweeps import ( ADMIN_ONLY_ALLOWANCES, ALLOWANCE_SLOT_FAMILIES, @@ -43,6 +38,12 @@ from integration.security._sweeps import ( sweep_routes, sweep_sink, ) +from redis import Redis + +from tests.integration._support.database import read_rows, write_rows +from tests.integration._support.redis_process import owned_redis +from tests.integration._support.tls import server_context, write_self_signed_cert +from tests.integration._support.workers import worker_services @pytest.fixture(scope="module") @@ -84,6 +85,8 @@ def test_rig_with_an_overridden_master_key_resolves_the_config_deployment(tmp_pa assert overridden.proxy.key == master_key assert overridden.model_id assert overridden.proxy.request("GET", "/model/info").status_code == 200 + assert overridden.proxy.request("GET", "/public/whats_new").status_code == 200 + assert overridden.egress() == () def _skills_upstream(request: Request) -> Reply: diff --git a/tests/integration/translation/runner.py b/tests/integration/translation/runner.py index 427c55e9cc8..8e6587c916d 100644 --- a/tests/integration/translation/runner.py +++ b/tests/integration/translation/runner.py @@ -1,12 +1,34 @@ +import base64 import json +from collections.abc import Mapping from typing import Final +from pydantic import JsonValue + from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration._support.wire import Reply from integration.translation.case import TranslationTestCase +from litellm.responses.utils import ResponsesAPIRequestUtils -TRANSPORT_HEADERS: Final = frozenset({"host", "accept", "accept-encoding", "connection", "content-length", "user-agent"}) +TRANSPORT_HEADERS: Final = frozenset( + {"host", "accept", "accept-encoding", "connection", "content-length", "user-agent"} +) + + +def expected_response_for_deployment( + expected: Mapping[str, JsonValue], model_id: str | None +) -> Mapping[str, JsonValue]: + identifier: Final = expected.get("id") + if not isinstance(identifier, str) or not ResponsesAPIRequestUtils._is_litellm_encoded_response_id(identifier): + return expected + assert model_id, "An encoded response ID must identify the selected deployment" + provider, _, response = base64.b64decode(identifier.removeprefix("resp_")).decode().split(";", 2) + updated: Final = f"{provider};model_id:{model_id};{response}" + return { + **expected, + "id": "resp_" + base64.b64encode(updated.encode()).decode(), + } def assert_translation(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: @@ -22,4 +44,16 @@ def assert_translation(case: TranslationTestCase, gateway: Gateway, provider: Sh } == dict(case.expected_provider_headers) assert json.loads(sent.body) == case.expected_provider_request assert response.status_code == case.expected_litellm_status_code, response.text - assert response.json() == case.expected_litellm_response + model_id: Final = response.headers.get("x-litellm-model-id") + expected: Final = expected_response_for_deployment(case.expected_litellm_response, model_id) + if expected is not case.expected_litellm_response: + deployments: Final = gateway.get("/model/info")["data"] + assert isinstance(deployments, list), deployments + assert any( + isinstance(deployment, dict) + and deployment.get("model_name") == case.litellm_request["model"] + and isinstance(deployment.get("model_info"), dict) + and deployment["model_info"].get("id") == model_id + for deployment in deployments + ), "The encoded response ID must match the configured deployment" + assert response.json() == expected diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index 50c750de759..b65c9fa7ec6 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -3,21 +3,93 @@ from __future__ import annotations import asyncio import errno import importlib +import os import socket +import time from collections.abc import Iterator from pathlib import Path from types import ModuleType from typing import Final +from unittest.mock import MagicMock import httpx +import psutil import pytest +import yaml +from pydantic import TypeAdapter from tests.integration._support.client import Gateway +from tests.integration._support.process import worker_configuration TESTS_DIR: Final = Path(__file__).resolve().parents[2] UNRELATED_CRASH: Final = "Traceback (most recent call last):\nModuleNotFoundError: No module named 'litellm'\n" +def test_worker_configuration_routes_stock_edges_to_its_own_servers_and_preserves_other_targets( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1") + monkeypatch.setenv("INTEGRATION_PROVIDER_PORT", "8203") + monkeypatch.setenv("INTEGRATION_UPSTREAM_URL", "http://127.0.0.1:8202") + source: Final = tmp_path / "source.yaml" + original: Final = ( + 'api_base: http://127.0.0.1:8191/v1\ntoken_url: "http://127.0.0.1:8190/_oauth/token"\n' + "peer: http://127.0.0.1:81910/v1\ncustom: http://127.0.0.1:9000/v1\n" + ) + source.write_text(original) + configured: Final = worker_configuration(source, tmp_path / "worker") + assert yaml.safe_load(configured.read_text()) == { + "api_base": "http://127.0.0.1:8203/v1", + "token_url": "http://127.0.0.1:8202/_oauth/token", + "peer": "http://127.0.0.1:81910/v1", + "custom": "http://127.0.0.1:9000/v1", + "general_settings": {"database_connection_pool_limit": 2}, + } + assert source.read_text() == original + + +@pytest.mark.parametrize("limit", (1, 5, 0, None, "", "bad")) +def test_worker_configuration_preserves_explicit_pool_limits( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, limit: int | str | None +) -> None: + monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1") + source: Final = tmp_path / "source.yaml" + original: Final = { + "general_settings": {"database_connection_pool_limit": limit, "database_connection_timeout": 3}, + "model_list": [{"model_name": "synthetic", "litellm_params": {"model": "openai/synthetic"}}], + } + source.write_text(yaml.safe_dump(original)) + configured: Final = worker_configuration(source, tmp_path / "worker") + assert yaml.safe_load(configured.read_text()) == original + assert yaml.safe_load(source.read_text()) == original + + +def test_worker_configuration_budgets_implicit_pools_without_changing_connection_options( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1") + source: Final = tmp_path / "source.yaml" + original: Final = {"general_settings": {"database_extra_connection_params": {"connection_limit": "1"}}} + source.write_text(yaml.safe_dump(original)) + configured: Final = worker_configuration(source, tmp_path / "worker") + decoded: Final = TypeAdapter(dict[str, object]).validate_python(yaml.safe_load(configured.read_text())) + assert decoded == { + "general_settings": { + "database_connection_pool_limit": 2, + "database_extra_connection_params": {"connection_limit": "1"}, + } + } + assert yaml.safe_load(source.read_text()) == original + + +def test_serial_proxy_keeps_the_requested_configuration(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.delenv("INTEGRATION_WORKER_ISOLATION", raising=False) + source: Final = tmp_path / "source.yaml" + source.write_text("api_base: http://127.0.0.1:8191/v1\n") + assert worker_configuration(source, tmp_path / "unused") == source + assert source.read_text() == "api_base: http://127.0.0.1:8191/v1\n" + + @pytest.fixture def process_module(monkeypatch: pytest.MonkeyPatch) -> ModuleType: monkeypatch.syspath_prepend(str(TESTS_DIR)) @@ -64,6 +136,94 @@ def test_lost_port_race_needs_the_process_to_have_exited( assert not process_module._lost_port_race(None, _written_log(tmp_path, bind_error_line)) +@pytest.mark.parametrize("failure", (psutil.AccessDenied, psutil.NoSuchProcess)) +def test_an_unreadable_process_cannot_establish_listener_ownership( + process_module: ModuleType, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, failure: type[psutil.Error] +) -> None: + member: Final = MagicMock(spec=psutil.Process) + member.pid = os.getpid() + member.net_connections.side_effect = failure(member.pid) + monkeypatch.setattr(psutil, "process_iter", lambda: iter((member,))) + parent: Final = MagicMock() + parent.pid = os.getpgid(member.pid) + launch: Final = process_module._Launch(parent, 8123, tmp_path / "unused.log", 0.0) + assert not process_module._owns_listener(launch) + + +def test_an_unreadable_member_does_not_hide_another_owned_listener( + process_module: ModuleType, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + unreadable: Final = MagicMock(spec=psutil.Process) + unreadable.pid = 1 + unreadable.net_connections.side_effect = psutil.AccessDenied(1) + readable: Final = MagicMock(spec=psutil.Process) + readable.pid = 2 + listener: Final = MagicMock() + listener.status = psutil.CONN_LISTEN + listener.laddr.port = 8123 + readable.net_connections.return_value = (listener,) + monkeypatch.setattr(psutil, "process_iter", lambda: iter((unreadable, readable))) + monkeypatch.setattr(os, "getpgid", lambda pid: 99) + parent: Final = MagicMock() + parent.pid = 99 + launch: Final = process_module._Launch(parent, 8123, tmp_path / "unused.log", 0.0) + assert process_module._owns_listener(launch) + + +@pytest.mark.parametrize( + ("status", "ready"), + ( + (psutil.STATUS_RUNNING, True), + (psutil.STATUS_SLEEPING, True), + (psutil.STATUS_ZOMBIE, False), + (psutil.STATUS_DEAD, False), + (psutil.STATUS_STOPPED, False), + (psutil.STATUS_TRACING_STOP, False), + (psutil.AccessDenied, False), + (psutil.NoSuchProcess, False), + ), +) +def test_only_live_readable_workers_can_satisfy_readiness( + process_module: ModuleType, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, status: str | type[psutil.Error], ready: bool +) -> None: + member: Final = MagicMock(spec=psutil.Process) + member.pid = os.getpid() + if isinstance(status, str): + member.status.return_value = status + else: + member.status.side_effect = status(member.pid) + listener: Final = MagicMock() + listener.status = psutil.CONN_LISTEN + listener.laddr.port = 8123 + member.net_connections.return_value = (listener,) + monkeypatch.setattr(psutil, "process_iter", lambda: iter((member,))) + parent: Final = MagicMock() + parent.pid = os.getpgid(member.pid) + parent.poll.return_value = None + launch: Final = process_module._Launch(parent, 8123, tmp_path / "unused.log", 0.0) + original_client: Final = httpx.Client + + 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( + lambda request: httpx.Response(200, headers={"x-integration-worker-pid": str(member.pid)}) + ), + ) + + monkeypatch.setattr(httpx, "Client", client) + monkeypatch.setattr(time, "monotonic", lambda: 0.0) + monkeypatch.setenv("INTEGRATION_PROXY_READY_SECONDS", "-1") + if ready: + process_module._wait_until_ready(launch) + else: + with pytest.raises(AssertionError, match="readiness deadline"): + process_module._wait_until_ready(launch) + + @pytest.mark.parametrize( ("writer", "explicit_reader", "expected_reader"), ( @@ -110,5 +270,11 @@ def test_owned_runtime_configuration_preserves_custom_worker_startup_hooks( ) expected: Final = "custom:configure" if explicit else ("" if remove else "ambient:configure") assert environment["LITELLM_WORKER_STARTUP_HOOKS"] == ",".join( - hook for hook in ("integration._support.runtime:configure_executor", expected) if hook + hook + for hook in ( + "integration._support.runtime:configure_executor", + "integration._support.runtime:configure_readiness", + expected, + ) + if hook ) diff --git a/tests/unit/integration_support/test_translation.py b/tests/unit/integration_support/test_translation.py new file mode 100644 index 00000000000..fad08043e85 --- /dev/null +++ b/tests/unit/integration_support/test_translation.py @@ -0,0 +1,34 @@ +import base64 +from pathlib import Path +from typing import Final + +import pytest +from pydantic import JsonValue + + +@pytest.mark.parametrize("response_id", ("resp_original", "", "resp_original;part", "resp_response_id:original")) +def test_response_id_uses_the_deployment_id_and_preserves_provider_identity_and_the_complete_body( + response_id: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2])) + from integration.translation.runner import expected_response_for_deployment + + original: Final = f"litellm:custom_llm_provider:openai;model_id:recorded-deployment;response_id:{response_id}" + expected: Final[dict[str, JsonValue]] = { + "id": "resp_" + base64.b64encode(original.encode()).decode(), + "content": [{"text": "answer"}], + } + result: Final = expected_response_for_deployment(expected, "isolated-deployment") + updated: Final = f"litellm:custom_llm_provider:openai;model_id:isolated-deployment;response_id:{response_id}" + assert result == {**expected, "id": "resp_" + base64.b64encode(updated.encode()).decode()} + assert expected["id"] == "resp_" + base64.b64encode(original.encode()).decode() + with pytest.raises(AssertionError, match="must identify the selected deployment"): + expected_response_for_deployment(expected, None) + + +def test_unencoded_provider_identity_is_preserved_exactly(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2])) + from integration.translation.runner import expected_response_for_deployment + + expected: Final[dict[str, JsonValue]] = {"id": "resp_original", "content": "answer"} + assert expected_response_for_deployment(expected, "isolated-deployment") is expected diff --git a/tests/unit/integration_support/test_upstream.py b/tests/unit/integration_support/test_upstream.py new file mode 100644 index 00000000000..91498f0879c --- /dev/null +++ b/tests/unit/integration_support/test_upstream.py @@ -0,0 +1,52 @@ +from collections import deque +from pathlib import Path +from typing import Final + +import pytest +from starlette.types import Message, Scope +from starlette.websockets import WebSocket + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", (OSError("client disconnected"), RuntimeError("unexpected upstream failure"))) +async def test_realtime_disconnect_after_session_creation_keeps_unexpected_errors_visible( + failure: Exception, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[2])) + from integration._support.upstream import Provider + from integration.cost_calculation.cost_tracking_case import RealtimeResponse + + provider: Final = Provider() + response: Final = RealtimeResponse(content_type="application/x-realtime", events=({"type": "response.done"},)) + provider.scenario_store.put("owned-scenario", response) + incoming: Final[deque[Message]] = deque( + ({"type": "websocket.connect"}, {"type": "websocket.receive", "text": '{"type":"response.create"}'}) + ) + outgoing: Final[deque[Message]] = deque() + + async def receive() -> Message: + return incoming.popleft() + + async def send(message: Message) -> None: + if message["type"] == "websocket.send" and any(item["type"] == "websocket.send" for item in outgoing): + raise failure + outgoing.append(message) + + scope: Final[Scope] = { + "type": "websocket", + "path": "/v1/realtime", + "headers": [(b"authorization", b"Bearer owned-scenario")], + "query_string": b"model=synthetic", + } + socket: Final = WebSocket(scope, receive, send) + if isinstance(failure, RuntimeError): + with pytest.raises(RuntimeError, match="unexpected upstream failure"): + await provider.realtime(socket) + else: + await provider.realtime(socket) + assert [item["type"] for item in outgoing] == ["websocket.accept", "websocket.send"] + observed: Final = tuple(provider.observations.get_nowait() for _ in range(provider.observations.qsize())) + assert [(item.method, item.body) for item in observed] == [ + ("WEBSOCKET", {"query": [["model", "synthetic"]]}), + ("WEBSOCKET_FRAME", {"type": "response.create"}), + ] diff --git a/tests/unit/integration_support/test_workers.py b/tests/unit/integration_support/test_workers.py new file mode 100644 index 00000000000..457b953e481 --- /dev/null +++ b/tests/unit/integration_support/test_workers.py @@ -0,0 +1,51 @@ +from dataclasses import dataclass +from typing import Final + +import httpx +import pytest + +from tests.integration._support.client import Gateway +from tests.integration._support.workers import readiness_failure + +HEALTHY: Final = { + "status": "healthy", + "cache_type": "redis", + "ping_response": True, + "set_cache_response": "success", +} + + +@dataclass(frozen=True, slots=True) +class Subscribers: + count: int + + def pubsub_numsub(self, *channels: str) -> tuple[tuple[str, int], ...]: + return tuple((channel, self.count) for channel in channels) + + +@pytest.mark.parametrize("subscribers", (0, 1, 2)) +def test_readiness_requires_a_subscriber_for_each_proxy(subscribers: int) -> None: + def respond(request: httpx.Request) -> httpx.Response: + if request.url.path == "/cache/ping": + assert request.headers["authorization"] == "Bearer owned-key" + return httpx.Response(200, json=HEALTHY) + return httpx.Response(200) + + with httpx.Client(base_url="http://127.0.0.1:4000", transport=httpx.MockTransport(respond)) as client: + primary: Final = Gateway(client, "owned-key", "http://127.0.0.1:8190") + failure: Final = readiness_failure(client, primary.upstream_url, (primary, primary), Subscribers(subscribers)) + assert failure == (None if subscribers == 2 else f"auth-cache subscribers {subscribers}/2") + + +@pytest.mark.parametrize( + ("field", "value"), + (("status", "unhealthy"), ("cache_type", "local"), ("ping_response", False), ("set_cache_response", "failure")), +) +def test_readiness_rejects_a_proxy_without_working_shared_redis(field: str, value: str | bool) -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={**HEALTHY, field: value}) + + with httpx.Client(base_url="http://127.0.0.1:4000", transport=httpx.MockTransport(respond)) as client: + primary: Final = Gateway(client, "owned-key", "http://127.0.0.1:8190") + failure: Final = readiness_failure(client, primary.upstream_url, (primary,), Subscribers(1)) + assert failure == "proxy 0 Redis read/write health"