fix(ci): isolate integration workers and repair test fixtures (#45863)

* test(integration): wait for customer spend rows and use owned OpenRouter cost

* ci: isolate integration workers and clean up resources after crashes

* docs: describe isolated extensions workers and shard census

* fix(ci): restore integration service and fixture guarantees

* fix(integration): preserve replay IDs and native route fixtures

* fix(tests): scope integration harness imports to each unit case

* fix(ci): preserve fixtures and wait for owned workers

* fix(ci): isolate nested runs and telemetry cases

* fix(ci): make worker readiness and fault tests deterministic

* chore(ci): drop MCP fixture repair already covered by main

* chore(ci): keep the local catalog fix from main

---------

Co-authored-by: yuneng <yuneng@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-10-11 01:56:02 -07:00 • committed by GitHub
parent e0398455d0
commit 3090f4b23e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
38 changed files with 1476 additions and 240 deletions

View file

@ -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

View file

@ -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

View file

@ -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__":

View file

@ -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<void> {
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);

View file

@ -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/<provider>.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 `/<scenario_id>`, 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

View file

@ -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()

View file

@ -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))

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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)

View file

@ -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):

View file

@ -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,
}
)

View file

@ -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)

View file

@ -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

View file

@ -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

View file

@ -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())

View file

@ -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

View file

@ -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")

View file

@ -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"))

View file

@ -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:

View file

@ -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:

View file

@ -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:

View file

@ -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

View file

@ -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")

View file

@ -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

View file

@ -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

View file

@ -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(

View file

@ -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:

View file

@ -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:

View file

@ -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],

View file

@ -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,

View file

@ -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:

View file

@ -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

View file

@ -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
)

View file

@ -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

View file

@ -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"}),
]

View file

@ -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"