scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

Record each e2e test's steps from the harness it calls

Browse source 
A test's JUnit report says whether it passed, never what it did or where a failing test died. This records that from the harness, so nothing about it is hand-written and it cannot drift from what the test actually ran

`@step("create team with a budget")` from the new tests/e2e/e2e_metadata.py goes on harness helpers, never on tests, and appends its label to the running test's step log in call order. The label is recorded before the wrapped call, so a helper that raises still leaves its own label last: a failing test's last step is where it died. Every public harness method that performs an action now carries one, 355 across the client modules, lifecycle, idp, the logging readers, migrations and the claude_code driver

Only the outermost step records, tracked per thread. Harness layers call each other (ResourceManager.key goes through ProxyClient.generate_key, a domain client wraps the shared ProxyClient), so every layer carries a label and the story still reads at the level the test called in at, one beat per action. A step above @contextmanager holds the guard through __enter__ and __exit__, so a context's cleanup never lands behind the step a test died on, and a bare generator function is refused at import because its body interleaves with its caller's. Consecutive duplicates collapse and the log caps at 50, so a poll loop is one beat rather than fifty. The wrapper is a frame, so the eight cleanup and retry warnings raised directly inside decorated helpers use stacklevel=2 + STEP_FRAMES to keep reporting at their caller

The log is emptied first thing in pytest_runtest_setup and attached from the existing pytest_runtest_makereport wrapper after setup and again after call, so a test that errors in a fixture keeps the steps recorded before the crash. Teardown does not attach: finalizer steps are cleanup. Each attach drops the item's earlier step entries, so the second attach and a --reruns 1 retry replace the story rather than doubling it

Steps ride out as repeated <property name="step"> entries behind the fixed package/covers/source prefix, which stays byte-identical. The project-releaser emitter already regroups them into the results JSON's steps array. test_junit_report.py runs real pytest with --junitxml against this conftest, in-process and under -n 2, and pins the passing, failing, setup-error, rerun and wide-scope-fixture cases on the parsed XML
This commit is contained in:
ryan-crabbe-berri 2026-09-21 18:48:59 -07:00
parent d8627938aa
commit d0e37d39c4
44 changed files with 1233 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
tests/e2e/AGENTS.md
View file

@ -130,6 +130,14 @@ Current limits: Bedrock cannot be mounted in record or replay (SigV4 signs the H
The harness is fully typed with no error budget: `make lint-e2e-basedpyright` must report zero basedpyright errors, and CI enforces that on any PR touching `tests/e2e/**/*.py`. When a response field is untyped, model it in `models.py` (just the fields you read) and let pydantic validate it, rather than threading a `dict` or `Any` through the test
## Recorded test steps
`@step("POST /chat/completions")` from `e2e_metadata.py` goes on HARNESS helpers - client methods, `ResourceManager.key`, poll loops - never on a test, and appends its label to the running test's `user_properties` in call order. The list IS the test's user story, and because the label is recorded BEFORE the wrapped call, a failing test's LAST step is where it died. Nothing about steps is hand-written: the call sequence cannot drift from what the test actually did. A new public harness method that performs an action (an HTTP call, a poll, a login, a CLI run) gets a `@step`; pure builders, parsers and `_private` helpers do not. Labels are static, lowercase, one beat of the story: a plain-English action ("create team with a budget") or, for a raw route call, the route itself ("POST /v1/messages")
Only the outermost step records. Harness layers call each other - `ResourceManager.key` goes through `ProxyClient.generate_key`, a domain client wraps the shared `ProxyClient` - so every layer carries its own label and the story still reads at the level the test called in at, one beat per action. On a `@contextmanager` helper `@step` goes ABOVE `@contextmanager`: the setup and cleanup around its `yield` count as part of the step, while the `with` body - the test's own code - records as usual, so cleanup never lands behind the step a test died on. A bare generator function is refused at import, since its body interleaves with the caller's. A decorated helper that warns about its caller uses `stacklevel=2 + STEP_FRAMES`, because the wrapper is a frame too. Nesting is tracked per thread, so a helper that fans work out to worker threads still records their steps. Consecutive duplicates collapse and the log caps at 50 entries, so a poll loop is one beat rather than fifty. The log is emptied first thing in every test's setup phase and attached after setup and again after call, so a test that errors in a fixture keeps the steps recorded before the crash. Teardown steps are left out on purpose: they are cleanup, and listing them would put a finalizer's step after the one a failing test died on
Steps ride out as repeated JUnit `<property name="step">` entries (`junit_properties.py`), one per step rather than one delimiter-joined value, since a free-text label has no separator that can be reserved. The results JSON downstream regroups them into a `steps` array. `test_junit_report.py` runs real pytest with `--junitxml`, in-process and under `-n 2`, and pins what reaches the XML
## Coverage registry
The set of tests we want is a registry checked into this repo, one row per behavior; that file is the definition of done and the denominator. Each e2e test declares what it covers with `@pytest.mark.covers("...")`, and a small collector diffs the registry against the tests and ships coverage to the existing Grafana. No Allure, no new dependencies

9
tests/e2e/a2a/a2a_client.py
View file

@ -19,6 +19,7 @@ from pydantic import BaseModel, ConfigDict, Field
from e2e_config import settle_propagation
from e2e_http import NoBody, Result, Success, get_external, is_ok
from e2e_metadata import STEP_FRAMES, step
from proxy_client import ProxyClient
@ -291,6 +292,7 @@ class A2AResponse(BaseModel):
class A2AClient:
proxy: ProxyClient
@step("register A2A agent")
def register_agent(self, body: AgentRegisterBody) -> Result[AgentResponse]:
"""Register an agent and, on success, wait until the data plane serves it.
@ -337,6 +339,7 @@ class A2AClient:
)
time.sleep(self.proxy.poll_interval)
@step("get A2A agent")
def get_agent(self, agent_id: str) -> Result[AgentResponse]:
return self.proxy.transport.get(
f"/v1/agents/{agent_id}",
@ -345,6 +348,7 @@ class A2AClient:
response_type=AgentResponse,
)
@step("delete A2A agent")
def delete_agent(self, agent_id: str) -> None:
result = self.proxy.transport.delete(
f"/v1/agents/{agent_id}",
@ -353,8 +357,9 @@ class A2AClient:
response_type=NoBody,
)
if not is_ok(result):
warnings.warn(f"delete_agent({agent_id!r}) failed: {result}", stacklevel=2)
warnings.warn(f"delete_agent({agent_id!r}) failed: {result}", stacklevel=2 + STEP_FRAMES)
@step("GET /a2a/{id}/.well-known/agent-card.json")
def agent_card(self, agent_id: str, key: str) -> Result[ServedAgentCard]:
return self.proxy.transport.get(
f"/a2a/{agent_id}/.well-known/agent-card.json",
@ -363,6 +368,7 @@ class A2AClient:
response_type=ServedAgentCard,
)
@step("POST /a2a/{id} (message/send)")
def send_message(self, agent_id: str, key: str, body: A2AJsonRpcRequest) -> Result[A2AResponse]:
return self.proxy.transport.post(
f"/a2a/{agent_id}",
@ -376,6 +382,7 @@ def build_a2a_client(proxy: ProxyClient) -> A2AClient:
return A2AClient(proxy=proxy)
@step("fetch the published upstream agent card")
def fetch_agent_card(url: str, *, timeout: float = 20.0) -> Result[UpstreamAgentCard]:
"""Fetch a live A2A agent card from its /.well-known endpoint and parse it into the
registration model, so a test can register a real published card verbatim rather

10
tests/e2e/access_control/access_control_client.py
View file

@ -9,6 +9,7 @@ from pydantic import BaseModel, ValidationError
from proxy_client import ProxyClient
from e2e_http import NoBody, StreamingResponse, is_ok, unwrap
from e2e_metadata import step
from models import (
ChatBody,
ChatMessage,
@ -59,14 +60,17 @@ def error_envelope(body: str) -> ApiErrorEnvelope | None:
class AccessControlClient:
proxy: ProxyClient
@step("generate virtual key for LLM routes only")
def llm_only_key(self) -> str:
return self.proxy.generate_key(
KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"])
)
@step("delete virtual key")
def delete_key(self, key: str) -> None:
self.proxy.delete_key(key)
@step("POST /chat/completions")
def chat_status(
self, key: str, model: str, content: str, max_completion_tokens: int | None = None
) -> StreamingResponse:
@ -80,6 +84,7 @@ class AccessControlClient:
),
)
@step("create team")
def create_team(self, team_alias: str, models: list[str]) -> str:
team_id = unwrap(
self.proxy.transport.post(
@ -92,6 +97,7 @@ class AccessControlClient:
self._await_team(team_id)
return team_id
@step("set the team's model allow-list")
def set_team_models(self, team_id: str, team_alias: str, models: list[str]) -> None:
"""Replace the team's allow-list. /model/new appends a team-scoped deployment's
public name to it, so a test that means to grant only an access group has to
@ -105,6 +111,7 @@ class AccessControlClient:
)
)
@step("delete team")
def delete_team(self, team_id: str) -> None:
_ = self.proxy.transport.post(
"/team/delete",
@ -113,6 +120,7 @@ class AccessControlClient:
response_type=NoBody,
)
@step("read the model access group's info")
def access_group_info(self, access_group: str) -> AccessGroupInfoResponse | None:
result = self.proxy.transport.get(
f"/access_group/{access_group}/info",
@ -122,6 +130,7 @@ class AccessControlClient:
)
return unwrap(result) if is_ok(result) else None
@step("read the team's models")
def team_models(self, team_id: str) -> list[str] | None:
result = self.proxy.transport.get(
"/team/info",
@ -139,6 +148,7 @@ class AccessControlClient:
time.sleep(self.proxy.poll_interval)
raise AssertionError(f"/team/info never resolved team {team_id!r} created by /team/new")
@step("POST /model/new")
def create_model_status(self, key: str, model_name: str) -> StreamingResponse:
return self.proxy.transport.send(
"/model/new",

3
tests/e2e/batches/batch_cleanup.py
View file

@ -7,6 +7,7 @@ from typing import Final, Protocol
from batch_client import BatchObject, FileDeleteResponse
from capabilities import is_cloud_storage_id, is_managed_id
from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError
from e2e_metadata import step
from pydantic import BaseModel
CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0)
@ -50,6 +51,7 @@ def _require_cleanup_success[R: BaseModel](result: Result[R], operation: str) ->
raise AssertionError(f"{operation} failed: {result.kind}")
@step("delete uploaded file")
def cleanup_file(client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None = None) -> None:
delete: Final[Callable[[], Result[FileDeleteResponse]]] = (
(lambda: client.delete_file_as_admin(file_id, provider=provider))
@ -65,6 +67,7 @@ def cleanup_file(client: BatchCleanupClient, file_id: str, *, key: str, provider
), f"Delete file {file_id} did not confirm deletion"
@step("cancel batch and wait for a terminal status")
def cleanup_batch(
client: BatchCleanupClient,
batch_id: str,

12
tests/e2e/batches/batch_client.py
View file

@ -25,6 +25,7 @@ from e2e_http import (
StreamingResponse,
UnknownApiError,
)
from e2e_metadata import step
from models import LiteLLMParamsBody
UPLOAD_FILENAME = "batch_input.jsonl"
@ -136,12 +137,15 @@ def is_result_access_denied[R: BaseModel](result: Result[R]) -> bool:
class BatchClient:
proxy: ProxyClient
@step("register batch deployment")
def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str:
return self.proxy.create_model(model_name, litellm_params, mode="batch")
@step("delete batch deployment")
def delete_model(self, model_id: str) -> None:
self.proxy.delete_model(model_id)
@step("POST /v1/files")
def upload_file(
self,
*,
@ -161,6 +165,7 @@ class BatchClient:
response_type=FileObject,
)
@step("GET /v1/files/{id}")
def retrieve_file(
self, file_id: str, *, key: str, provider: str | None = None
) -> Result[FileObject]:
@ -171,6 +176,7 @@ class BatchClient:
response_type=FileObject,
)
@step("GET /v1/files")
def list_files(self, *, key: str, provider: str | None = None) -> Result[FileList]:
return self.proxy.transport.get(
_files_path(provider),
@ -179,6 +185,7 @@ class BatchClient:
response_type=FileList,
)
@step("POST /v1/batches")
def create_batch(
self, *, body: BatchCreateBody, key: str, provider: str | None = None
) -> StreamingResponse:
@ -188,6 +195,7 @@ class BatchClient:
json=body,
)
@step("GET /v1/batches/{id}")
def retrieve_batch(
self, batch_id: str, *, key: str, provider: str | None = None
) -> Result[BatchObject]:
@ -198,6 +206,7 @@ class BatchClient:
response_type=BatchObject,
)
@step("POST /v1/batches/{id}/cancel")
def cancel_batch(
self, batch_id: str, *, key: str, provider: str | None = None
) -> Result[BatchObject]:
@ -208,6 +217,7 @@ class BatchClient:
response_type=BatchObject,
)
@step("GET /v1/batches")
def list_batches(
self,
*,
@ -223,6 +233,7 @@ class BatchClient:
response_type=BatchList,
)
@step("DELETE /v1/files/{id}")
def delete_file(
self, file_id: str, *, key: str, provider: str | None = None
) -> Result[FileDeleteResponse]:
@ -233,6 +244,7 @@ class BatchClient:
response_type=FileDeleteResponse,
)
@step("DELETE /v1/files/{id} (admin)")
def delete_file_as_admin(self, file_id: str, *, provider: str | None = None) -> Result[FileDeleteResponse]:
return self.proxy.transport.delete(
f"{_files_path(provider)}/{file_id}",

3
tests/e2e/claude_code/_basic_messaging.py
View file

@ -31,6 +31,8 @@ from typing import Any, Callable, Mapping, Sequence
import pytest
from e2e_metadata import step
from claude_code._env import require_proxy
from claude_code.cli_driver import (
ClaudeCLIError,
@ -74,6 +76,7 @@ def _count_stream_event_deltas(events: Sequence[Mapping[str, Any]]) -> int:
return count
@step("send a basic message via the claude CLI")
def run_basic_messaging_cell(
*,
compat_result,

3
tests/e2e/claude_code/_passthrough.py
View file

@ -58,6 +58,8 @@ from typing import Any, Callable, Dict, Mapping, Optional, Sequence
import pytest
from e2e_metadata import step
from claude_code._env import require_proxy
from claude_code.cli_driver import (
ClaudeCLIError,
@ -118,6 +120,7 @@ def foundry_extra_env(proxy_base_url: str) -> Dict[str, str]:
}
@step("run the claude CLI via a passthrough route")
def run_passthrough_cell(
*,
compat_result,

4
tests/e2e/claude_code/cli_driver.py
View file

@ -25,6 +25,8 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
from e2e_metadata import step
from claude_code.rate_limiter import (
RateLimiter,
get_default_limiter,
@ -211,6 +213,7 @@ class DriverResult:
duration_ms: Optional[int] = None
@step("run the claude CLI")
def run_claude(
*,
prompt: Optional[str],
@ -394,6 +397,7 @@ def _matches_failure_shape(outcome: ModelResult, pattern: "re.Pattern[str]") ->
return bool(pattern.search(failure_diagnostic(outcome)))
@step("run the claude CLI")
def run_claude_models_parallel(
*,
models: Sequence[str],

4
tests/e2e/claude_code/http_probe.py
View file

@ -42,6 +42,7 @@ from e2e_http import (
UnknownApiError,
ValidationError,
)
from e2e_metadata import step
from models import (
AnthropicAssistantTurn,
AnthropicCustomTool,
@ -109,6 +110,7 @@ def _acquire(model: str, rate_limiter: RateLimiter | None) -> None:
limiter.acquire(infer_provider(model))
@step("POST /v1/messages/count_tokens")
def probe_count_tokens(
*,
client: ProxyClient,
@ -132,6 +134,7 @@ def probe_count_tokens(
)
@step("POST /v1/messages with the tool_search tool")
def probe_tool_search(
*,
client: ProxyClient,
@ -228,6 +231,7 @@ def _replay_history(answer: AnthropicMessagesResponse) -> tuple[AnthropicMessage
)
@step("replay a tool_search turn as history")
def probe_tool_search_multiturn(
*,
client: ProxyClient,

33
tests/e2e/conftest.py
View file

@ -39,10 +39,11 @@ from e2e_config import (
)
from e2e_db import RESET_OPT_IN_ENV, reset_spend_logs, run_spend_log_cleanup
from e2e_http import unwrap
from e2e_metadata import STEPS
from fixture_mode import fixture_mode_collection_error, fixture_report_lines
from fixture_mode import pytest_fixture_setup as pytest_fixture_setup
from idp import Identity, Keycloak, keycloak_from_env
from junit_properties import attach_result_properties
from junit_properties import attach_result_properties, attach_step_properties
from lifecycle import ProxyClientProvider, ResourceManager
from models import TeamNewBody, UserNewBody, UserNewResponse
from provider_cache_routing import LIVE_PROVIDER_REQUIRED
@ -232,7 +233,14 @@ def pytest_runtest_setup(item: pytest.Item) -> None:
"""Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe.
Unmarked tests (unit coverage of the harness) don't touch the proxy, so they
run even when none is up. Never skip for a missing proxy. Replay mode needs
the proxy too: only provider-bound traffic replays from the bundle."""
the proxy too: only provider-bound traffic replays from the bundle.
Also empties the step log, so the story a test tells is its own. It happens
here, first in the setup phase, rather than in a fixture: a fixture only runs
once every wider-scoped fixture ahead of it has been set up, so a step a
module-scoped finalizer recorded after the previous test would still be in
the log when this test's setup dies early, and would be reported as its own."""
STEPS.reset()
LIVE_PROVIDER_REQUIRED.set(item.get_closest_marker("provider_live") is not None)
if item.get_closest_marker("e2e") is None or item.get_closest_marker("migration_startup") is not None:
return
@ -259,7 +267,24 @@ def pytest_runtest_makereport(
item: pytest.Item, call: pytest.CallInfo[None]
) -> Generator[None, pytest.TestReport, pytest.TestReport]:
"""Stash the call-phase outcome so teardown can tell a passed test from a
failed one without re-deriving it."""
failed one without re-deriving it, and attach the runtime-recorded steps.
The steps cannot ride along with the other properties in
`pytest_collection_modifyitems`: that hook runs before any test body has, so
the recorder is empty there. They are attached after setup and again after
call, on every outcome -- a failing test's last step is where it died, which
is the whole reason the field exists. Setup has to attach too because a test
whose fixture raises never reaches the call phase, and setup is where an e2e
test most often dies (proxy not ready, key creation failing). The second
attach replaces the first, so nothing is doubled. JUnit writes properties
from the teardown report, which pytest builds from `item.user_properties`
after both of these have run.
Teardown deliberately does not attach. Steps recorded by fixture finalizers
are cleanup, and appending them would put "delete virtual key" after the step
a failing test died on, which breaks the one guarantee the field makes. A
finalizer that raises is still reported by JUnit with its own traceback.
"""
report = yield
if item.get_closest_marker("mcp_oauth_live") is not None and call.excinfo is not None:
# Publish code locations only, never exception messages, source text or locals.
@ -270,6 +295,8 @@ def pytest_runtest_makereport(
report.user_properties = list(item.user_properties)
if report.when == "call":
item.stash[_CALL_PASSED] = report.passed
if report.when in ("setup", "call"):
attach_step_properties(item)
return report

183
tests/e2e/e2e_metadata.py Normal file
View file

@ -0,0 +1,183 @@
"""Per-test metadata for the e2e suite: the step log each test records as it runs.
`steps` is appended at runtime by `@step`-decorated harness helpers, in call
order, so the list IS the test's user story and its last element is where a
failing test died. Nothing about it is hand-written, so it cannot drift from
what the test actually did.
Stdlib-only on purpose. tests/e2e is a black-box HTTP suite that imports litellm
in zero files and is shipped to the runner image as tests/e2e alone, and every
harness module imports this one.
"""
from __future__ import annotations
import inspect
import threading
from collections.abc import Callable, Generator
from contextlib import AbstractContextManager, contextmanager
from functools import wraps
from types import TracebackType
from typing import Final, ParamSpec, TypeVar, cast
_P = ParamSpec("_P")
_R = TypeVar("_R")
_Y = TypeVar("_Y")
MAX_STEPS: Final = 50
MAX_STEP_CHARS: Final = 200
STEP_FRAMES: Final = 1
"""Frames a `@step` wrapper puts between a helper and its caller. A decorated
helper that warns about its caller adds this to `stacklevel`
(`stacklevel=2 + STEP_FRAMES`), or the warning is reported at the wrapper."""
class _StepRecorder:
"""The ordered step log for the running test.
A plain lock-guarded list rather than a ContextVar: ContextVars do not
propagate into worker threads, and several e2e helpers call out from
threads. Under xdist each worker is its own process, so there is no
cross-test bleed beyond what the per-test reset already handles.
"""
def __init__(self) -> None:
self._lock = threading.Lock()
self._steps: list[str] = []
def reset(self) -> None:
"""Called first thing in every test's setup phase, so each test starts
empty."""
with self._lock:
self._steps.clear()
def record(self, label: str) -> None:
"""Append `label`, unless it repeats the previous step or the log is full.
A retrying helper (poll_cost_row) or a load test calling a decorated
helper in a loop would otherwise emit thousands of <property> entries per
testcase: a consecutive repeat collapses, so a poll loop is one step in
the story rather than fifty, and the log stops growing at MAX_STEPS.
"""
cleaned = " ".join(label.split())[:MAX_STEP_CHARS]
if not cleaned:
return
with self._lock:
if self._steps and self._steps[-1] == cleaned:
return
if len(self._steps) >= MAX_STEPS:
return
self._steps.append(cleaned)
def taken(self) -> tuple[str, ...]:
with self._lock:
return tuple(self._steps)
STEPS: Final = _StepRecorder()
class _Nesting(threading.local):
"""Whether this thread is already inside a `@step` helper.
Per thread, like the helpers themselves: a worker thread a step fans out to
starts outside any step, so its own decorated calls still record."""
def __init__(self) -> None:
self.inside: bool = False
_NESTING: Final = _Nesting()
@contextmanager
def _inside_step() -> Generator[None]:
"""Hold the nesting guard for the duration, restoring whatever it was."""
outer: Final = _NESTING.inside
_NESTING.inside = True
try:
yield
finally:
_NESTING.inside = outer
class _StepContext(AbstractContextManager[_Y]):
"""A `@contextmanager` helper's context, entered and exited inside its step.
Calling a `@contextmanager` function runs none of its body: the setup runs at
`__enter__` and the cleanup at `__exit__`, both after the call has returned
and so both outside the guard the call held. Here each runs inside it, so the
helpers they call stay out of the story, while the `with` body in between --
the test's own code -- still records. Without this, a test that died inside
the `with` would have the cleanup's steps appended behind the one it died on.
"""
def __init__(self, inner: AbstractContextManager[_Y]) -> None:
self._inner: Final = inner
def __enter__(self) -> _Y:
with _inside_step():
return self._inner.__enter__()
def __exit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
traceback: TracebackType | None,
) -> bool | None:
with _inside_step():
return self._inner.__exit__(exc_type, exc, traceback)
def step(label: str) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
"""Record `label` on the running test whenever this helper is called.
Goes on HARNESS helpers (client methods, fixtures), never on tests. The
label is recorded BEFORE the wrapped call, so a helper that raises still
leaves its own label as the last element -- which is the whole point: the
last step is where the test died.
Only the outermost step records. Harness layers call each other --
`ResourceManager.key` goes through `ProxyClient.generate_key`, a domain
client wraps the shared `ProxyClient` -- so every layer can carry its own
label without one action showing up in the story once per layer. The story
reads at the level the test called in at, and the label of the helper the
test called is still the last one when anything beneath it raises.
On a `@contextmanager` helper `@step` goes ABOVE `@contextmanager`, and the
setup and cleanup around its `yield` count as part of the step (see
`_StepContext`). A bare generator function is refused where the decorator
runs: its body only runs as the caller iterates, interleaved with the
caller's own steps, so no single point in the story is where it happened.
"""
def decorate(fn: Callable[_P, _R]) -> Callable[_P, _R]:
if inspect.isgeneratorfunction(fn):
raise TypeError(
f"@step({label!r}) cannot wrap the generator function {fn!r}: put it on a helper that"
" returns, or above @contextmanager on one that yields a context"
)
underlying: Final[object] = inspect.unwrap(fn) # pyright: ignore[reportAny] # inspect.unwrap is typed as returning Any
opens_a_context: Final = inspect.isgeneratorfunction(underlying)
@wraps(fn)
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
if not _NESTING.inside:
STEPS.record(label)
with _inside_step():
result = fn(*args, **kwargs)
if opens_a_context and isinstance(result, AbstractContextManager):
context: Final = cast("AbstractContextManager[object]", result)
return cast("_R", _StepContext(context))
return result
return wrapper
return decorate
def step_properties() -> tuple[tuple[str, str], ...]:
"""The step log as repeated `step` properties. Appended after the setup and
call phases, never at collection."""
return tuple(("step", label) for label in STEPS.taken())

20
tests/e2e/guardrails/guardrails_client.py
View file

@ -11,6 +11,7 @@ from typing import Final, Literal
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation, unique_marker
from e2e_http import NoBody, Result, StreamingResponse, Success, unwrap
from e2e_metadata import step
from lifecycle import ResourceManager
from models import (
AnthropicMessagesBody,
@ -154,6 +155,7 @@ class _ResponsesGuardrailBody(BaseModel):
class GuardrailsClient:
proxy: ProxyClient
@step("register content-filter guardrail")
def create_content_filter_guardrail(self, name: str, blocked_keyword: str, *, default_on: bool = True) -> str:
return self.register(
name,
@ -164,6 +166,7 @@ class GuardrailsClient:
),
)
@step("register Bedrock guardrail")
def create_bedrock_guardrail(
self,
name: str,
@ -191,6 +194,7 @@ class GuardrailsClient:
),
)
@step("register deployment")
def create_backend_model(
self,
resources: ResourceManager,
@ -211,6 +215,7 @@ class GuardrailsClient:
resources.defer(lambda: self.proxy.delete_model(model_id))
return model_name
@step("register guardrail")
def register(self, name: str, params: GuardrailParamsBody) -> str:
"""Register any guardrail via POST /guardrails and return its id, once every
replica can be expected to serve it. New built-ins register with
@ -235,6 +240,7 @@ class GuardrailsClient:
settle_propagation(time.monotonic())
return guardrail_id
@step("delete guardrail")
def delete_guardrail(self, guardrail_id: str) -> None:
_ = self.proxy.transport.delete(
f"/guardrails/{guardrail_id}",
@ -243,6 +249,7 @@ class GuardrailsClient:
response_type=NoBody,
)
@step("create team opted out of global guardrails")
def create_team_opted_out_of_global_guardrails(self, alias: str) -> str:
team_id = unwrap(
self.proxy.transport.post(
@ -258,6 +265,7 @@ class GuardrailsClient:
self._await_team(team_id)
return team_id
@step("delete team")
def delete_team(self, team_id: str) -> None:
_ = self.proxy.transport.post(
"/team/delete",
@ -266,9 +274,11 @@ class GuardrailsClient:
response_type=NoBody,
)
@step("generate virtual key in the team")
def create_key_in_team(self, team_id: str) -> str:
return self.proxy.generate_key(KeyGenerateBody(team_id=team_id, user_id="e2e-guardrails-user"))
@step("generate virtual key with guardrails")
def create_key_with_guardrails(self, resources: ResourceManager, guardrails: list[str]) -> str:
key = self.proxy.generate_key(
KeyGenerateBody(user_id="e2e-guardrails-user", metadata=KeyMetadata(guardrails=guardrails))
@ -276,6 +286,7 @@ class GuardrailsClient:
resources.defer(lambda: self.proxy.delete_key(key))
return key
@step("POST /v1/videos")
def create_video(self, key: str, model: str, prompt: str) -> Result[VideoCreateResponse]:
return self.proxy.transport.post(
"/v1/videos",
@ -284,6 +295,7 @@ class GuardrailsClient:
response_type=VideoCreateResponse,
)
@step("POST /chat/completions")
def chat(
self,
key: str,
@ -310,6 +322,7 @@ class GuardrailsClient:
),
)
@step("POST /chat/completions")
def chat_raw(
self,
key: str,
@ -338,6 +351,7 @@ class GuardrailsClient:
),
)
@step("POST /chat/completions (streaming)")
def chat_stream_raw(
self,
key: str,
@ -362,6 +376,7 @@ class GuardrailsClient:
),
)
@step("POST /v1/messages")
def messages(
self,
key: str,
@ -381,6 +396,7 @@ class GuardrailsClient:
),
)
@step("POST /v1/responses")
def responses(
self,
key: str,
@ -395,6 +411,7 @@ class GuardrailsClient:
json=_ResponsesGuardrailBody(model=model, input=text, guardrails=guardrails),
)
@step("POST /guardrails/apply_guardrail")
def apply_guardrail(self, key: str, *, name: str, text: str) -> Result[ApplyGuardrailResponse]:
return self.proxy.transport.post(
"/guardrails/apply_guardrail",
@ -423,6 +440,7 @@ def build_client(proxy: ProxyClient) -> GuardrailsClient:
return GuardrailsClient(proxy=proxy)
@step("poll the request for the applied guardrail")
def poll_until_guardrail_applied(
call: Callable[[], StreamingResponse],
guardrail_name: str,
@ -446,6 +464,7 @@ def poll_until_guardrail_applied(
return result
@step("poll the request for a guardrail block")
def poll_until_blocked[R: BaseModel](call: Callable[[], Result[R]]) -> Result[R]:
"""Retry a call that a guardrail should reject until it is, returning the last result.
@ -473,6 +492,7 @@ def poll_until_blocked[R: BaseModel](call: Callable[[], Result[R]]) -> Result[R]
_TRANSIENT_STREAM_STATUSES = frozenset({-1, 401, 429})
@step("poll the request for a guardrail block")
def poll_until_blocked_stream(call: Callable[[], StreamingResponse]) -> StreamingResponse:
"""poll_until_blocked for raw/streamed sends, which return a StreamingResponse
instead of a Result: retry while the call still succeeds (the data-plane worker

16
tests/e2e/idp.py
View file

@ -31,6 +31,7 @@ from e2e_http import (
post_json_external,
unwrap,
)
from e2e_metadata import step
from pydantic import BaseModel, Field
KEYCLOAK_URL_ENV: Final = "E2E_KEYCLOAK_URL"
@ -171,6 +172,7 @@ class Keycloak:
case _:
return pytest.fail(f"Keycloak refused {context}: {result}")
@step("create Keycloak group")
def create_group(self, name: str) -> str:
return created_id(
post_json_external(
@ -179,6 +181,7 @@ class Keycloak:
f"group {name}",
)
@step("create Keycloak user")
def create_user(
self, *, username: str, email: str, password: str, group: str | None = None, groups: tuple[str, ...] = ()
) -> str:
@ -196,12 +199,15 @@ class Keycloak:
f"user {username}",
)
@step("delete Keycloak user")
def delete_user(self, user_id: str) -> None:
self._delete(f"/users/{user_id}")
@step("delete Keycloak group")
def delete_group(self, group_id: str) -> None:
self._delete(f"/groups/{group_id}")
@step("check the Keycloak resource is gone")
def assert_absent(self, kind: Literal["users", "groups", "clients"], resource_id: str) -> None:
result: Final = get_external(
self._admin_url(f"/{kind}/{resource_id}"),
@ -230,11 +236,13 @@ class Keycloak:
stacklevel=2,
)
@step("provision a Keycloak user in a new group")
def provision(self, *, marker: str, group: str, defer: Callable[[Callable[[], object]], None]) -> Identity:
"""Create `group` and a user in it, credentialed with a password generated
for this test alone, and hand back the identity a token can be minted for."""
return self.provision_groups(marker=marker, groups=(group,), defer=defer)
@step("provision a Keycloak user in new groups")
def provision_groups(
self, *, marker: str, groups: tuple[str, ...], defer: Callable[[Callable[[], object]], None]
) -> Identity:
@ -246,6 +254,7 @@ class Keycloak:
group_ids: Final = tuple(provision_group(group) for group in groups)
return self.provision_user(marker=marker, groups=groups, group_ids=group_ids, defer=defer)
@step("provision a Keycloak user in existing groups")
def provision_user(
self,
*,
@ -262,6 +271,7 @@ class Keycloak:
defer(lambda: self.delete_user(user_id))
return Identity(user_id=user_id, username=username, password=password, groups=groups, group_ids=group_ids)
@step("sign in to Keycloak for an access token")
def access_token(
self, identity: Identity, *, client_id: str = TESTS_CLIENT_ID, issuer_host: str | None = None
) -> str:
@ -275,9 +285,11 @@ class Keycloak:
)
return self._token(result, f"a token for {identity.username}")
@step("read Keycloak's OIDC discovery document")
def discovery(self) -> Discovery:
return unwrap(get_external(f"{self.issuer}/.well-known/openid-configuration", response_type=Discovery))
@step("register a Keycloak browser client")
def browser_client(self, *, callback_url: str, defer: Callable[[Callable[[], object]], None]) -> BrowserClient:
client: Final = BrowserClient(
client_id=f"e2e-browser-{secrets.token_hex(8)}",
@ -309,6 +321,7 @@ class Keycloak:
assert configured.attributes.pkce == "S256"
return client
@step("sign in to Keycloak via the browser client")
def browser_token(self, identity: Identity, client: BrowserClient) -> str:
return self._token(
post_form_external(
@ -325,6 +338,7 @@ class Keycloak:
"browser-profile identity mapping",
)
@step("read Keycloak userinfo for the token")
def userinfo(self, token: str) -> UserInfo:
return unwrap(
get_external(
@ -435,6 +449,7 @@ def _signal_process_group(process_id: int, signum: int) -> bool:
return True
@step("stop the child process group")
def stop_process_group(child: subprocess.Popen[bytes]) -> None:
_signal_process_group(child.pid, signal.SIGTERM)
deadline: Final = time.monotonic() + 5
@ -457,6 +472,7 @@ def _process_group_exists(process_id: int) -> bool:
return True
@step("run a command under the OIDC browser profile")
def run_oidc_profile(proxy_url: str, command: list[str]) -> int:
idp: Final = keycloak_from_env().with_strict_cleanup()
with ExitStack() as cleanup:

19
tests/e2e/junit_properties.py
View file

@ -20,6 +20,7 @@ from collections.abc import Iterable
import pytest
from coverage_registry.management_cases import case_properties
from e2e_metadata import step_properties
# Hardcoded because the runner image copies tests/e2e/ to /app/e2e, so nothing
# at runtime names this suite's place in the repo. test_junit_properties.py
@ -105,3 +106,21 @@ def attach_result_properties(item: pytest.Item) -> None:
if any(name == "package" for name, _ in item.user_properties):
return
item.user_properties.extend(result_properties(item))
def attach_step_properties(item: pytest.Item) -> None:
"""Attach the runtime-recorded steps; called after setup and after call.
Separate from `attach_result_properties` because it cannot share its home:
that one runs in `pytest_collection_modifyitems`, before any test body has
executed, so the recorder is necessarily empty there.
Any `step` entries already on the item are dropped first, which is what makes
the second call of a test safe: the story attached after setup is replaced by
the longer one attached after call. It also covers `--reruns 1`, where a flaky
test's second attempt would otherwise append a second copy of the story behind
the first, and the report would read as one very long test that did everything
twice. Last attempt wins, which is the attempt whose outcome JUnit records.
"""
item.user_properties[:] = [entry for entry in item.user_properties if entry[0] != "step"]
item.user_properties.extend(step_properties())

3
tests/e2e/lifecycle.py
View file

@ -12,6 +12,7 @@ from builtins import ExceptionGroup
from dataclasses import dataclass, field
from typing import Callable, Final, List, Protocol, runtime_checkable
from e2e_metadata import step
from proxy_client import ProxyClient
from models import KeyGenerateBody
@ -69,6 +70,7 @@ class ResourceManager:
response model can be deferred directly."""
self._cleanups.append(cleanup)
@step("generate virtual key")
def key(self, models: list[str] | None = None, user_id: str | None = "e2e-test-user") -> str:
"""Create a virtual key; delete it on teardown. `models` restricts which
models the key may call (None/[] means all). `user_id` is required for
@ -78,6 +80,7 @@ class ResourceManager:
self.defer(lambda: self.client.delete_key(key))
return key
@step("register end user")
def customer(self, customer_id: str) -> str:
"""Track an end-user id (from the `user` param); delete it on teardown."""
self.defer(lambda: self.client.delete_customers([customer_id]))

12
tests/e2e/llm_translation/passthrough_client.py
View file

@ -20,6 +20,7 @@ from websockets.sync.client import connect
from e2e_config import ws_base_url
from proxy_client import ProxyClient
from e2e_http import FileUploadForm, Headers, NoBody, Result, StreamingResponse
from e2e_metadata import step
from models import ChatMessage
@ -246,6 +247,7 @@ class PassthroughClient:
# ---- Gemini native passthrough (/gemini/v1beta/...) -----------------
@step("call Gemini generateContent via /gemini")
def gemini_generate(
self,
key: str,
@ -263,6 +265,7 @@ class PassthroughClient:
),
)
@step("call Gemini streamGenerateContent via /gemini")
def gemini_stream(
self, key: str, model: str, text: str, *, tags: list[str] | None = None
) -> StreamingResponse:
@ -278,6 +281,7 @@ class PassthroughClient:
# ---- Vertex AI native passthrough (/vertex_ai/v1/projects/...) -------
@step("call Vertex generateContent via /vertex_ai")
def vertex_generate(
self, key: str, project: str, location: str, model: str, text: str
) -> StreamingResponse:
@ -295,6 +299,7 @@ class PassthroughClient:
# ---- Anthropic native passthrough (/anthropic/v1/messages) ----------
@step("POST /anthropic/v1/messages")
def anthropic_message(
self,
key: str,
@ -324,6 +329,7 @@ class PassthroughClient:
# Relayed to OpenAI untouched, which is the whole point of the prefix: the
# customer opts out of the gateway's managed-file handling here.
@step("POST /openai_passthrough/v1/files")
def openai_passthrough_upload_file(
self, key: str, *, content: bytes, filename: str
) -> Result[PassthroughFileObject]:
@ -336,6 +342,7 @@ class PassthroughClient:
response_type=PassthroughFileObject,
)
@step("DELETE /openai_passthrough/v1/files/{id}")
def openai_passthrough_delete_file(
self, key: str, file_id: str
) -> Result[PassthroughFileDeleted]:
@ -346,6 +353,7 @@ class PassthroughClient:
response_type=PassthroughFileDeleted,
)
@step("GET /openai_passthrough/v1/batches")
def openai_passthrough_list_batches(self, key: str) -> Result[PassthroughBatchList]:
return self.proxy.transport.get(
"/openai_passthrough/v1/batches",
@ -360,6 +368,7 @@ class PassthroughClient:
# budgets against this traffic, so a 200 that logs no spend is money the
# gateway never sees.
@step("POST /openai_passthrough/v1/responses")
def openai_passthrough_responses(
self, key: str, model: str, text: str, *, stream: bool = False
) -> StreamingResponse:
@ -370,6 +379,7 @@ class PassthroughClient:
stream=stream,
)
@step("POST /openai_passthrough/v1/embeddings")
def openai_passthrough_embed(
self, key: str, model: str, text: str
) -> StreamingResponse:
@ -379,6 +389,7 @@ class PassthroughClient:
json=OpenAIEmbeddingBody(model=model, input=text),
)
@step("POST /openai/v1/chat/completions")
def openai_chat(
self, key: str, model: str, text: str, *, max_completion_tokens: int = 64
) -> StreamingResponse:
@ -397,6 +408,7 @@ class PassthroughClient:
# The same prefixes over an upgrade instead of a POST, for the provider APIs
# that only speak websocket (realtime, responses.connect).
@step("open a passthrough websocket")
def openai_passthrough_websocket(
self,
key: str,

5
tests/e2e/llm_translation/realtime/realtime_client.py
View file

@ -23,6 +23,7 @@ from websockets.sync.connection import Connection
from e2e_config import unique_marker, ws_base_url
from proxy_client import ProxyClient
from e2e_metadata import step
from models import LiteLLMParamsBody
_M = TypeVar("_M", bound=BaseModel)
@ -294,9 +295,11 @@ def _as_text(message: str | bytes) -> str:
class RealtimeSession:
connection: Connection
@step("send a realtime client event")
def send(self, event: BaseModel) -> None:
self.connection.send(event.model_dump_json(by_alias=True, exclude_none=True))
@step("poll the realtime socket for the awaited event")
def collect_until(
self, stop_type: str, *, timeout: float
) -> tuple[ReceivedEvent, ...]:
@ -324,6 +327,7 @@ class RealtimeSession:
class RealtimeClient:
proxy: ProxyClient
@step("register a realtime deployment")
def provision(self, provider: RealtimeProvider) -> tuple[str, str]:
"""Register this provider's realtime deployment through /model/new and return
(model_name, model_id). The name is marker-unique so it never collides with a
@ -336,6 +340,7 @@ class RealtimeClient:
)
return model_name, model_id
@step("open the /v1/realtime websocket")
@contextmanager
def connect(
self, *, key: str, model: str, timeout: float = 15.0

3
tests/e2e/load/session_anomaly.py
View file

@ -9,6 +9,7 @@ from pydantic import BaseModel
from e2e_config import unique_marker
from e2e_http import Result, Success
from e2e_metadata import step
from models import CacheControl, RichMessage, TextBlock
from transport import Transport
@ -213,6 +214,7 @@ def _drive_turns(
)
@step("run a cached /v1/messages session")
def run_session(
transport: Transport, key: str, model: str, turns: int, attempts_per_turn: int
) -> tuple[TurnMetric, ...]:
@ -230,6 +232,7 @@ def run_session(
)
@step("run concurrent cached /v1/messages sessions")
def run_concurrent_sessions(
transport: Transport,
key: str,

5
tests/e2e/logging/datadog_reader.py
View file

@ -32,6 +32,7 @@ from e2e_config import (
POLL_TIMEOUT,
)
from e2e_http import URL, Headers, StreamingResponse, send
from e2e_metadata import step
type SearchCall = Callable[[str, float], StreamingResponse]
@ -115,6 +116,7 @@ class DdLogsReader:
sleep: Callable[[float], None] = field(default=time.sleep, repr=False)
jitter: Callable[[], float] = field(default=random.random, repr=False)
@step("search DataDog logs for the marker")
def events_for_marker(self, marker: str) -> list[DdLogEvent]:
"""Every ingested event whose attributes carry the marker. DataDog
consumes the shipped JSON message into ``attributes`` and leaves the
@ -124,6 +126,7 @@ class DdLogsReader:
it)."""
return self.events_for_query(f"*:*{marker}*")
@step("search DataDog logs for the query")
def events_for_query(self, query: str) -> list[DdLogEvent]:
"""Every ingested event the search query matches (failure payloads
carry no prompt to mark, so failure scenarios query indexed attributes
@ -156,10 +159,12 @@ class DdLogsReader:
timeout=timeout,
)
@step("poll DataDog logs for the marker")
def poll_events_for_marker(self, marker: str) -> list[DdLogEvent]:
"""``poll_events_for_query`` over the every-attribute marker scan."""
return self.poll_events_for_query(f"*:*{marker}*")
@step("poll DataDog logs for the query")
def poll_events_for_query(self, query: str) -> list[DdLogEvent]:
"""Poll until at least one matching event is searchable (the callback
flushes in periodic batches and DataDog ingestion adds seconds of lag),

3
tests/e2e/logging/gcs_reader.py
View file

@ -30,6 +30,7 @@ from pydantic import BaseModel, ConfigDict, Field
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT
from e2e_http import URL, Headers, probe
from e2e_metadata import step
_GCS_API = "https://storage.googleapis.com"
#: Tolerance for clock skew between this host and GCS object timestamps.
@ -146,6 +147,7 @@ class GcsLogReader:
)
return result.body
@step("read GCS log records for the response")
def records_for_response_id(self, response_id: str, *, since: datetime) -> list[GcsLogRecord]:
"""Every payload written for ``response_id``: the direct
``{date}/{response_id}`` object plus any hit inside batch NDJSON
@ -168,6 +170,7 @@ class GcsLogReader:
)
return records
@step("poll GCS for the response's log records")
def poll_records_for_response_id(self, response_id: str, *, since: datetime) -> list[GcsLogRecord]:
"""Poll until the payload is readable (the gcs_bucket callback flushes
on a ~20s timer), then keep re-reading for GCS_SETTLE_SECONDS - past a

25
tests/e2e/logging/logging_client.py
View file

@ -35,6 +35,7 @@ from e2e_http import (
get,
unwrap,
)
from e2e_metadata import step
from models import (
AnthropicMessagesBody,
ChatBody,
@ -303,6 +304,7 @@ def observation_has_guardrail(obs: LangfuseObservation, *, guardrail_name: str)
class LoggingClient:
proxy: ProxyClient
@step("generate virtual key")
def key_with_alias(
self,
alias: str,
@ -324,9 +326,11 @@ class LoggingClient:
)
)
@step("delete virtual key")
def delete_key(self, key: str) -> None:
self.proxy.delete_key(key)
@step("create team")
def create_team(
self,
alias: str,
@ -347,6 +351,7 @@ class LoggingClient:
)
).team_id
@step("delete team")
def delete_team(self, team_id: str) -> None:
_ = self.proxy.transport.post(
"/team/delete",
@ -355,6 +360,7 @@ class LoggingClient:
response_type=NoBody,
)
@step("create internal user")
def create_user(self, *, user_email: str, user_id: str | None = None) -> str:
return unwrap(
self.proxy.transport.post(
@ -369,6 +375,7 @@ class LoggingClient:
)
).user_id
@step("delete internal user")
def delete_user(self, user_id: str) -> None:
_ = self.proxy.transport.post(
"/user/delete",
@ -377,6 +384,7 @@ class LoggingClient:
response_type=NoBody,
)
@step("create organization")
def create_org(self, alias: str, *, models: list[str]) -> str:
return unwrap(
self.proxy.transport.post(
@ -387,6 +395,7 @@ class LoggingClient:
)
).organization_id
@step("delete organization")
def delete_org(self, organization_id: str) -> None:
_ = self.proxy.transport.delete(
"/organization/delete",
@ -395,6 +404,7 @@ class LoggingClient:
response_type=NoBody,
)
@step("add a Langfuse callback to the team")
def add_team_langfuse_callback(
self,
team_id: str,
@ -418,6 +428,7 @@ class LoggingClient:
f"POST /team/{team_id}/callback must return status=success; got {response.status!r}"
)
@step("create tool-permission guardrail")
def create_tool_permission_guardrail(self, name: str, *, allowed_tool: str) -> str:
"""Register a tool_permission guardrail that allows one tool and denies the rest."""
response = unwrap(
@ -451,6 +462,7 @@ class LoggingClient:
settle_propagation(time.monotonic())
return guardrail_id
@step("delete guardrail")
def delete_guardrail(self, guardrail_id: str) -> None:
_ = self.proxy.transport.delete(
f"/guardrails/{guardrail_id}",
@ -459,12 +471,15 @@ class LoggingClient:
response_type=NoBody,
)
@step("register deployment")
def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str:
return self.proxy.create_model(model_name, litellm_params)
@step("delete deployment")
def delete_model(self, model_id: str) -> None:
self.proxy.delete_model(model_id)
@step("POST /chat/completions")
def chat(self, key: str, model: str, text: str) -> ChatResponse:
return unwrap(
self.proxy.chat(
@ -477,6 +492,7 @@ class LoggingClient:
)
)
@step("POST /chat/completions")
def chat_raw(
self,
key: str,
@ -506,6 +522,7 @@ class LoggingClient:
json=body,
)
@step("POST /v1/messages")
def messages_raw(
self, key: str, model: str, text: str, *, max_tokens: int = 16, stream: bool = False
) -> StreamingResponse:
@ -522,6 +539,7 @@ class LoggingClient:
return self.proxy.transport.stream("/v1/messages", headers=self.proxy.transport.bearer(key), json=body)
return self.proxy.transport.send("/v1/messages", headers=self.proxy.transport.bearer(key), json=body)
@step("POST /v1/responses")
def responses_raw(
self, key: str, model: str, text: str, *, max_output_tokens: int = 64, stream: bool = False
) -> StreamingResponse:
@ -537,9 +555,11 @@ class LoggingClient:
return self.proxy.transport.stream("/v1/responses", headers=self.proxy.transport.bearer(key), json=body)
return self.proxy.transport.send("/v1/responses", headers=self.proxy.transport.bearer(key), json=body)
@step("GET /metrics")
def scrape_metrics(self) -> str:
return self.proxy.probe("/metrics", params=NoBody()).body
@step("poll /spend/logs for the key")
def poll_proxy_spend_for_key(
self,
key: str,
@ -567,6 +587,7 @@ class LoggingClient:
return row
return None
@step("list Langfuse observations")
def list_langfuse_observations(
self,
creds: LangfuseCreds,
@ -593,6 +614,7 @@ class LoggingClient:
case _:
return []
@step("find the run's Langfuse observation")
def find_langfuse_observation(
self,
creds: LangfuseCreds,
@ -611,6 +633,7 @@ class LoggingClient:
return obs
return None
@step("poll Langfuse for the run's observation")
def poll_langfuse_observation(
self,
creds: LangfuseCreds,
@ -630,6 +653,7 @@ class LoggingClient:
time.sleep(POLL_INTERVAL)
return last
@step("poll Langfuse for the run's trace")
def poll_langfuse_trace_observations(
self,
creds: LangfuseCreds,
@ -663,6 +687,7 @@ def build_logging_client(proxy: ProxyClient) -> LoggingClient:
return LoggingClient(proxy=proxy)
@step("GET /health/readiness/details")
def readiness_details_body(client: LoggingClient) -> str:
"""/health/readiness/details, tolerating the 503 it serves while the
ephemeral stack's DB leg blips: the recorded state the logging suites check

5
tests/e2e/logging/s3_reader.py
View file

@ -24,6 +24,7 @@ import pytest
from pydantic import BaseModel, ConfigDict
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT
from e2e_metadata import step
if TYPE_CHECKING:
from types_boto3_s3.client import S3Client
@ -53,17 +54,21 @@ class S3LogReader:
bucket: str
client: S3Client
@step("list S3 log objects")
def list_keys(self, prefix: str) -> list[str]:
response = self.client.list_objects_v2(Bucket=self.bucket, Prefix=prefix)
return [obj["Key"] for obj in response.get("Contents", []) if "Key" in obj]
@step("read S3 log record")
def read_record(self, key: str) -> S3LogRecord:
body = self.client.get_object(Bucket=self.bucket, Key=key)["Body"].read()
return S3LogRecord.model_validate_json(body)
@step("read matching S3 log records")
def records_matching(self, *, prefix: str, predicate: Callable[[S3LogRecord], bool]) -> list[S3LogRecord]:
return [record for record in map(self.read_record, self.list_keys(prefix)) if predicate(record)]
@step("poll S3 for matching log records")
def poll_records(self, *, prefix: str, predicate: Callable[[S3LogRecord], bool]) -> list[S3LogRecord]:
"""Poll until at least one matching object is listed (the s3_v2
callback flushes on a ~10s timer), then keep re-reading for

3
tests/e2e/logging/weave_reader.py
View file

@ -37,6 +37,7 @@ from pydantic import BaseModel, ConfigDict, Field
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT
from e2e_http import URL, AuthHeaders, send
from e2e_metadata import step
_WEAVE_TRACE_API: Final = "https://trace.wandb.ai"
@ -232,6 +233,7 @@ class WeaveReader:
)
return tuple(WeaveCall.model_validate_json(line) for line in outcome.body.splitlines() if line.strip())
@step("query Weave calls for the marker")
def calls_matching(self, marker: str, *, since: float, op: str = LITELLM_REQUEST_OP) -> tuple[WeaveCall, ...]:
"""Every call under ``op`` started after ``since`` whose inputs carry
``marker``, paging until the window is exhausted.
@ -247,6 +249,7 @@ class WeaveReader:
)
return tuple(call for page in pages for call in page if call.mentions(marker))
@step("poll Weave for the marker's calls")
def poll_calls_matching(self, marker: str, *, since: float, op: str = LITELLM_REQUEST_OP) -> tuple[WeaveCall, ...]:
"""Poll until the call is readable, then keep re-reading for
WEAVE_SETTLE_SECONDS so a duplicate exported by a later batch flush

5
tests/e2e/management/jwt_actors.py
View file

@ -5,6 +5,7 @@ from typing import Final, Literal
from e2e_config import unique_marker
from e2e_http import NoBody, unwrap
from e2e_metadata import step
from idp import ADMIN_CLIENT_ID, TESTS_CLIENT_ID, Identity, Keycloak
from lifecycle import ResourceManager
from management.management_client import ManagementClient
@ -53,6 +54,7 @@ class Actor:
profile: ActorProfile
tenants: tuple[Tenant, ...]
@step("log in to Keycloak as the actor")
def mint_caller(self, idp: Keycloak) -> Caller:
return Caller(
credential=idp.access_token(
@ -74,6 +76,7 @@ class ActorFactory:
if self.bootstrap.proxy.caller is not None:
raise ValueError("Actor bootstrap requires a separately held master client")
@step("generate virtual key")
def key(self, tenant: Tenant | None = None, *, user_id: str | None = None) -> KeyGenerateResponse:
created: Final = unwrap(
self.bootstrap.generate_key(
@ -87,6 +90,7 @@ class ActorFactory:
self.resources.defer(lambda: self.bootstrap.delete_key_strict(created.key, missing_ok=True))
return created
@step("create tenant org, team and Keycloak group")
def tenant(self) -> Tenant:
marker: Final = unique_marker()
organization_id: Final = self.bootstrap.create_org(OrgNewBody(organization_alias=f"e2e-organization-{marker}"))
@ -118,6 +122,7 @@ class ActorFactory:
self.resources.defer(lambda: self.idp.with_strict_cleanup().delete_group(group_id))
return Tenant(organization_id=organization_id, team_id=team_id, group_id=group_id)
@step("provision actor in Keycloak and the proxy")
def create(
self, role: ActorRole, *, tenants: tuple[Tenant, ...] = (), profile: ActorProfile = "database_role"
) -> Actor:

54
tests/e2e/management/management_client.py
View file

@ -24,6 +24,7 @@ from e2e_http import (
retry_attempts,
unwrap,
)
from e2e_metadata import STEP_FRAMES, step
from models import (
ChatBody,
ChatMessage,
@ -113,9 +114,11 @@ class ManagementClient:
def with_caller(self, caller: Caller) -> ManagementClient:
return replace(self, proxy=self.proxy.with_caller(caller))
@step("generate LLM-only virtual key")
def llm_only_key(self) -> str:
return self.proxy.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"]))
@step("generate virtual key")
def generate_key(self, body: KeyGenerateBody, *, caller_key: str | None = None) -> Result[KeyGenerateResponse]:
"""POST /key/generate. `caller_key` is who is creating the key: the master
key by default, or a virtual key (an admin filling in Create New Key on the
@ -130,6 +133,7 @@ class ManagementClient:
response_type=KeyGenerateResponse,
)
@step("update virtual key")
def update_key(self, body: KeyUpdateBody, *, caller_key: str | None = None) -> Result[NoBody]:
"""POST /key/update. `caller_key` is who is editing: the master key by
default, or a virtual key (the dashboard edits under the session key its
@ -149,16 +153,18 @@ class ManagementClient:
case UnknownApiError(body=error_body) if any(
marker in error_body.lower() for marker in _TRANSIENT_BACKEND_MARKERS
):
warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2)
warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2 + STEP_FRAMES)
time.sleep(0.5 * (attempt + 1))
continue
case _:
break
return last
@step("update virtual key's allowed models")
def update_key_models(self, key: str, models: list[str]) -> None:
_ = unwrap(self.update_key(KeyUpdateBody(key=key, models=models)))
@step("read /key/info")
def key_info_as(self, key: str, *, caller_key: str | None = None) -> Result[KeyInfoResponse]:
return self.proxy.transport.get(
"/key/info",
@ -167,6 +173,7 @@ class ManagementClient:
response_type=KeyInfoResponse,
)
@step("delete virtual key")
def delete_key_strict(self, key: str, *, caller_key: str | None = None, missing_ok: bool = False) -> None:
"""Strict delete for the act phase of a test: a failed delete is a hard
failure, unlike the warn-only ProxyClient.delete_key used at teardown."""
@ -180,6 +187,7 @@ class ManagementClient:
return
_ = unwrap(result)
@step("delete deployment")
def delete_model_strict(self, model_id: str) -> None:
"""Strict delete for the act phase of a test: a failed delete is a hard
failure, unlike the warn-only ProxyClient.delete_model used at teardown."""
@ -192,6 +200,7 @@ class ManagementClient:
)
)
@step("test connection to the provider")
def connection_test(self, body: ConnectionTestBody) -> Result[ConnectionTestResponse]:
"""POST /health/test_connection, the call behind the Admin UI's Test
Connection button, probing the live provider with the supplied params."""
@ -203,6 +212,7 @@ class ManagementClient:
timeout=120.0,
)
@step("block virtual key")
def block_key(self, key: str) -> None:
_ = unwrap(
self.proxy.transport.post(
@ -212,6 +222,7 @@ class ManagementClient:
response_type=NoBody,
)
)
@step("regenerate virtual key")
def regenerate_key(self, key: str, *, grace_period: str | None = None) -> str:
return unwrap(
self.proxy.transport.post(
@ -222,6 +233,7 @@ class ManagementClient:
)
).key
@step("reset virtual key spend")
def reset_key_spend(self, key: str, reset_to: float) -> KeyResetSpendResponse:
return unwrap(
self.proxy.transport.post(
@ -232,6 +244,7 @@ class ManagementClient:
)
)
@step("list virtual keys by alias")
def key_list(self, key_alias: str, *, caller_key: str | None = None) -> Result[KeyListResponse]:
"""GET /key/list, the Virtual Keys page's own inventory call. `caller_key` is
who is asking: the master key by default, or a virtual key."""
@ -243,9 +256,11 @@ class ManagementClient:
response_type=KeyListResponse,
)
@step("count virtual keys by alias")
def key_alias_count(self, key_alias: str) -> int:
return unwrap(self.key_list(key_alias)).total_count
@step("log in to the dashboard")
def dashboard_login(self, username: str, password: str) -> DashboardSession:
"""POST /v2/login, the call the Admin UI's sign-in form makes.
@ -269,6 +284,7 @@ class ManagementClient:
redirect_url=response.redirect_url,
)
@step("create team")
def create_team(self, body: TeamNewBody) -> str:
team_id = unwrap(
self.proxy.transport.post(
@ -281,6 +297,7 @@ class ManagementClient:
self._wait_for_team(team_id)
return team_id
@step("update team")
def update_team(self, body: TeamUpdateBody) -> None:
last: Result[NoBody] | None = None
for attempt in range(retry_attempts(5)):
@ -296,7 +313,7 @@ class ManagementClient:
case UnknownApiError(body=body_text) if (
"connecting to redis" in body_text.lower() or "name resolution" in body_text.lower()
):
warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2)
warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2 + STEP_FRAMES)
time.sleep(0.5 * (attempt + 1))
continue
case _:
@ -304,6 +321,7 @@ class ManagementClient:
assert last is not None
raise AssertionError(last)
@step("delete team")
def delete_team(self, team_id: str) -> None:
_ = self.proxy.transport.post(
"/team/delete",
@ -312,6 +330,7 @@ class ManagementClient:
response_type=NoBody,
)
@step("read /team/info")
def team_info(self, team_id: str) -> TeamData:
return unwrap(
self.proxy.transport.get(
@ -322,6 +341,7 @@ class ManagementClient:
)
).team_info
@step("list teams")
def team_list_ids(self) -> tuple[str, ...]:
return tuple(
entry.team_id
@ -335,6 +355,7 @@ class ManagementClient:
).root
)
@step("probe /team/info")
def team_info_status(self, team_id: str) -> ProbeResult:
return self.proxy.transport.probe(
"/team/info", params=TeamInfoParams(team_id=team_id), headers=self.proxy.management_headers()
@ -358,6 +379,7 @@ class ManagementClient:
assert last is not None
raise AssertionError(last)
@step("add a member to the team")
def add_team_member(self, team_id: str, user_id: str) -> None:
last: Result[NoBody] | None = None
for attempt in range(retry_attempts(_TEAM_READY_ATTEMPTS)):
@ -374,7 +396,7 @@ class ManagementClient:
_TEAM_READY_ATTEMPTS
):
warnings.warn(
"Retrying team membership while the team becomes available", RuntimeWarning, stacklevel=2
"Retrying team membership while the team becomes available", RuntimeWarning, stacklevel=2 + STEP_FRAMES
)
time.sleep(_TEAM_READY_SLEEP_SECONDS)
continue
@ -383,6 +405,7 @@ class ManagementClient:
assert last is not None
raise AssertionError(last)
@step("remove a member from the team")
def delete_team_member(self, team_id: str, user_id: str) -> None:
_ = unwrap(
self.proxy.transport.post(
@ -393,6 +416,7 @@ class ManagementClient:
)
)
@step("create internal user")
def create_user(self, body: UserNewBody) -> str:
return unwrap(
self.proxy.transport.post(
@ -403,6 +427,7 @@ class ManagementClient:
)
).user_id
@step("register end user")
def create_customer(self, user_id: str) -> str:
_ = unwrap(
self.proxy.transport.post(
@ -414,6 +439,7 @@ class ManagementClient:
)
return user_id
@step("read /customer/info")
def customer_info(self, end_user_id: str) -> CustomerResponse:
return unwrap(
self.proxy.transport.get(
@ -424,6 +450,7 @@ class ManagementClient:
)
)
@step("delete end user")
def delete_customer(self, user_id: str) -> None:
_ = self.proxy.transport.post(
"/customer/delete",
@ -432,6 +459,7 @@ class ManagementClient:
response_type=NoBody,
)
@step("update internal user")
def update_user(self, body: UserUpdateBody) -> None:
_ = unwrap(
self.proxy.transport.post(
@ -442,6 +470,7 @@ class ManagementClient:
)
)
@step("delete internal user")
def delete_user(self, user_id: str) -> None:
_ = self.proxy.transport.post(
"/user/delete",
@ -450,6 +479,7 @@ class ManagementClient:
response_type=NoBody,
)
@step("delete internal user")
def delete_user_strict(self, user_id: str) -> None:
"""Strict delete for the act phase of a test: a failed delete is a hard
failure, unlike the warn-only delete_user used at teardown."""
@ -462,6 +492,7 @@ class ManagementClient:
)
)
@step("read /user/info")
def user_info(self, user_id: str | None = None) -> UserInfoResponse:
return unwrap(
self.proxy.transport.get(
@ -472,6 +503,7 @@ class ManagementClient:
)
)
@step("count users in /user/list")
def user_count(self, user_id: str) -> int:
return unwrap(
self.proxy.transport.get(
@ -482,6 +514,7 @@ class ManagementClient:
)
).total
@step("list users in /user/list")
def user_list_ids(self, user_id: str) -> tuple[str, ...]:
listing = unwrap(
self.proxy.transport.get(
@ -493,6 +526,7 @@ class ManagementClient:
)
return tuple(row.user_id for row in listing.users)
@step("create organization")
def create_org(self, body: OrgNewBody) -> str:
return unwrap(
self.proxy.transport.post(
@ -503,6 +537,7 @@ class ManagementClient:
)
).organization_id
@step("update organization")
def update_org(self, body: OrgUpdateBody) -> None:
_ = unwrap(
self.proxy.transport.patch(
@ -513,6 +548,7 @@ class ManagementClient:
)
)
@step("delete organization")
def delete_org(self, organization_id: str) -> None:
_ = self.proxy.transport.delete(
"/organization/delete",
@ -521,6 +557,7 @@ class ManagementClient:
response_type=NoBody,
)
@step("read /organization/info")
def org_info(self, organization_id: str) -> OrgInfoResponse:
return unwrap(
self.proxy.transport.get(
@ -531,6 +568,7 @@ class ManagementClient:
)
)
@step("probe /organization/info")
def org_info_status(self, organization_id: str) -> ProbeResult:
return self.proxy.transport.probe(
"/organization/info",
@ -538,6 +576,7 @@ class ManagementClient:
headers=self.proxy.management_headers(),
)
@step("create tag")
def create_tag(self, body: TagNewBody) -> None:
_ = unwrap(
self.proxy.transport.post(
@ -548,6 +587,7 @@ class ManagementClient:
)
)
@step("delete tag")
def delete_tag(self, name: str) -> None:
_ = self.proxy.transport.post(
"/tag/delete",
@ -556,6 +596,7 @@ class ManagementClient:
response_type=NoBody,
)
@step("list tags")
def tag_list(self) -> tuple[TagListEntry, ...]:
return tuple(
unwrap(
@ -568,6 +609,7 @@ class ManagementClient:
).root
)
@step("create MCP server")
def create_mcp_server(self, body: McpServerCreateBody) -> McpServerRow:
return unwrap(
self.proxy.transport.post(
@ -578,6 +620,7 @@ class ManagementClient:
)
)
@step("update MCP server")
def update_mcp_server(self, body: McpServerUpdateBody) -> McpServerRow:
"""PUT /v1/mcp/server, the call behind the dashboard's Save Changes: a partial
update where a field left unset keeps its stored value and None clears it."""
@ -590,6 +633,7 @@ class ManagementClient:
)
)
@step("delete MCP server")
def delete_mcp_server(self, server_id: str) -> Result[NoBody]:
"""DELETE /v1/mcp/server/{server_id}. Returns the outcome so the act phase can
unwrap it while a deferred teardown can ignore an already-deleted server."""
@ -600,6 +644,7 @@ class ManagementClient:
response_type=NoBody,
)
@step("POST /chat/completions")
def chat_status(self, key: str, model: str, content: str) -> StreamingResponse:
return self.proxy.transport.send(
"/chat/completions",
@ -607,12 +652,15 @@ class ManagementClient:
json=ChatBody(model=model, messages=[ChatMessage(role="user", content=content)], max_tokens=16),
)
@step("POST /key/generate")
def key_generate_status(self, key: str, body: KeyGenerateBody) -> StreamingResponse:
return self.proxy.transport.send("/key/generate", headers=self.proxy.transport.bearer(key), json=body)
@step("POST /team/new")
def team_new_status(self, key: str, body: TeamNewBody) -> StreamingResponse:
return self.proxy.transport.send("/team/new", headers=self.proxy.transport.bearer(key), json=body)
@step("POST /user/new")
def user_new_status(self, key: str, body: UserNewBody) -> StreamingResponse:
return self.proxy.transport.send("/user/new", headers=self.proxy.transport.bearer(key), json=body)

2
tests/e2e/mcp/datadog_mcp.py
View file

@ -6,6 +6,7 @@ import os
from collections.abc import Sequence
from e2e_config import datadog_mcp_url, unique_marker
from e2e_metadata import step
from lifecycle import ResourceManager
from mcp_client import McpClient
@ -31,6 +32,7 @@ def assert_dd_mcp_creds() -> None:
)
@step("register the Datadog MCP server")
def register_datadog_mcp(
client: McpClient,
resources: ResourceManager,

16
tests/e2e/mcp/mcp_client.py
View file

@ -21,6 +21,7 @@ from pydantic import BaseModel, ConfigDict, Field, RootModel
from e2e_config import settle_propagation
from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap
from e2e_metadata import step
from models import KeyGenerateBody, McpServerListResponse, McpServerRow, ObjectPermission
from proxy_client import ProxyClient
@ -153,6 +154,7 @@ class McpCallToolResponse(BaseModel):
class McpClient:
proxy: ProxyClient
@step("register MCP server")
def register_server(
self,
*,
@ -183,6 +185,7 @@ class McpClient:
)
).server_id
@step("delete MCP server")
def delete_server(self, server_id: str) -> None:
_ = self.proxy.transport.delete(
f"/v1/mcp/server/{server_id}",
@ -191,6 +194,7 @@ class McpClient:
response_type=NoBody,
)
@step("list MCP servers as the admin")
def registered_servers(self) -> list[McpServerRow]:
return unwrap(
self.proxy.transport.get(
@ -201,6 +205,7 @@ class McpClient:
)
).root
@step("list MCP servers as the key")
def list_servers(self, key: str) -> Result[McpServerListResponse]:
return self.proxy.transport.get(
"/v1/mcp/server",
@ -209,6 +214,7 @@ class McpClient:
response_type=McpServerListResponse,
)
@step("check MCP server health")
def server_health(self, key: str, server_ids: list[str] | None = None) -> Result[McpHealthResponse]:
return self.proxy.transport.get(
"/v1/mcp/server/health",
@ -217,6 +223,7 @@ class McpClient:
response_type=McpHealthResponse,
)
@step("poll /v1/mcp/server for the server")
def await_registered(self, server_id: str) -> McpServerRow:
"""Wait for every configured replica to list the server and return its row."""
registered = self.proxy.read_body_back_everywhere(
@ -228,6 +235,7 @@ class McpClient:
row for response in registered.values() for row in response.root if row.server_id == server_id
)
@step("generate virtual key")
def generate_key(
self,
*,
@ -254,6 +262,7 @@ class McpClient:
)
)
@step("list MCP tools")
def list_tools(self, key: str) -> Result[McpToolsListResponse]:
return self.proxy.transport.get(
"/mcp-rest/tools/list",
@ -262,6 +271,7 @@ class McpClient:
response_type=McpToolsListResponse,
)
@step("poll /mcp-rest/tools/list for the tool")
def await_tool(self, key: str, server_id: str, needle: str) -> str:
"""Poll tools/list until `server_id` serves a tool matching `needle`, and
return its fully-qualified name. Fails at poll_timeout.
@ -287,6 +297,7 @@ class McpClient:
)
time.sleep(self.proxy.poll_interval)
@step("poll /mcp-rest/tools/list for the expected tools")
def await_tools(self, key: str, server_id: str, *, expected: frozenset[str]) -> frozenset[str]:
"""Poll tools/list until `server_id`'s tools as `key` sees them are exactly
`expected`, and return the last listing either way, so the caller's equality
@ -301,6 +312,7 @@ class McpClient:
return unwrap(result).tool_names_for_server(server_id)
time.sleep(self.proxy.poll_interval)
@step("poll /mcp-rest/tools/call for a result")
def await_call_tool(
self,
key: str,
@ -329,6 +341,7 @@ class McpClient:
)
time.sleep(self.proxy.poll_interval)
@step("poll /mcp-rest/tools/call for a 403 denial")
def await_call_tool_denied(
self,
key: str,
@ -355,6 +368,7 @@ class McpClient:
)
time.sleep(self.proxy.poll_interval)
@step("register MCP content-filter guardrail")
def register_mcp_content_filter(self, *, name: str, blocked_keyword: str) -> str:
"""Register a default-on content-filter guardrail that runs on the MCP
tool-call hook (pre_mcp_call) and blocks a single keyword. The keyword is
@ -378,6 +392,7 @@ class McpClient:
settle_propagation(time.monotonic())
return guardrail_id
@step("delete guardrail")
def delete_guardrail(self, guardrail_id: str) -> None:
_ = self.proxy.transport.delete(
f"/guardrails/{guardrail_id}",
@ -386,6 +401,7 @@ class McpClient:
response_type=NoBody,
)
@step("call MCP tool")
def call_tool(
self,
key: str,

9
tests/e2e/mcp/oauth_chat_client.py
View file

@ -26,6 +26,7 @@ import httpx2
import pytest
from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT
from e2e_http import AuthHeaders, NoBody, unwrap
from e2e_metadata import step
from idp import Identity
from mcp import ClientSession
from mcp.client.auth import OAuthClientProvider
@ -318,6 +319,7 @@ async def _list_and_call(
class ChatMcpClient:
proxy: ProxyClient
@step("register MCP server")
def create_server(self, body: McpServerCreateBody) -> McpServerInfo:
return unwrap(
self.proxy.transport.post(
@ -328,6 +330,7 @@ class ChatMcpClient:
)
)
@step("get MCP server info")
def server_info(self, server_id: str) -> McpServerInfo:
return unwrap(
self.proxy.transport.get(
@ -338,6 +341,7 @@ class ChatMcpClient:
)
)
@step("delete MCP server")
def delete_server(self, server_id: str) -> None:
_ = self.proxy.transport.delete(
f"/v1/mcp/server/{server_id}",
@ -346,6 +350,7 @@ class ChatMcpClient:
response_type=NoBody,
)
@step("seed upstream OAuth token via authorize dance")
def seed_user_token(self, alias: str, key: str, storage_state_path: str) -> tuple[str, ...]:
"""Drive the interactive authorize dance for `key`'s user so the gateway
stores their upstream token, retried to the shared deadline since the
@ -367,6 +372,7 @@ class ChatMcpClient:
f"last error: {last_error!r}"
)
@step("list and call an MCP tool via the SDK")
def list_and_call(
self,
alias: str,
@ -394,6 +400,7 @@ class ChatMcpClient:
)
)
@step("list MCP server user credentials")
def server_user_credentials(self, server_id: str) -> tuple[McpServerUserCredentialRow, ...]:
return unwrap(
self.proxy.transport.get(
@ -404,6 +411,7 @@ class ChatMcpClient:
)
).root
@step("revoke the user's MCP OAuth credential")
def revoke_user_token(self, server_id: str, headers: AuthHeaders) -> None:
_ = unwrap(
self.proxy.transport.delete(
@ -414,6 +422,7 @@ class ChatMcpClient:
)
)
@step("POST /chat/completions (with MCP tools)")
def chat_with_mcp(self, headers: AuthHeaders, body: ChatBody) -> ChatResponse:
"""POST /chat/completions carrying the LiteLLM key in `headers` (either
ingress form) with an MCP server attached in `body.tools`. The gateway

7
tests/e2e/mcp/oauth_gateway.py
View file

@ -21,6 +21,7 @@ from typing import Final
import psycopg
from e2e_http import NoBody
from e2e_metadata import step
from idp import Keycloak, stop_process_group
from proxy_client import ProxyClient, build_proxy_client
from psycopg.rows import class_row
@ -39,6 +40,7 @@ class CredentialRow:
credential_b64: str = field(repr=False)
@step("read stored OAuth credential from the DB")
def stored_oauth(user_id: str, server_id: str) -> StoredOAuth:
"""Read the encrypted credential because management APIs omit the plaintext token."""
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
@ -91,6 +93,7 @@ class OAuthObservation:
with self._lock:
self._seen = (*self._seen, (operation, received, gateway_leaked))
@step("assert upstream got the stored OAuth token")
def assert_forwarded(self, expected: StoredOAuth) -> None:
with self._lock:
snapshot: Final = self._seen
@ -116,6 +119,7 @@ class OAuthGateway:
_log_path: Path
_child: subprocess.Popen[bytes] | None = field(default=None, init=False, repr=False)
@step("start the owned OAuth gateway")
def start(self) -> None:
with self._log_path.open("ab") as log:
self._child = subprocess.Popen(
@ -134,11 +138,13 @@ class OAuthGateway:
time.sleep(0.5)
raise AssertionError("owned OAuth gateway did not become ready")
@step("stop the owned OAuth gateway")
def stop(self) -> None:
if self._child is not None:
stop_process_group(self._child)
assert self._child.poll() is not None, "old gateway process is still alive"
@step("restart the owned OAuth gateway")
def restart(self) -> None:
assert self._child is not None
previous: Final = self._child.pid
@ -147,6 +153,7 @@ class OAuthGateway:
assert self._child.pid != previous, "gateway restart did not create a new process"
@step("provision the owned OAuth gateway")
def owned_gateway(idp: Keycloak, directory: Path, cleanup: ExitStack) -> OAuthGateway:
for name in ("DATABASE_URL", "LITELLM_LICENSE", "LITELLM_SALT_KEY", "LITELLM_MASTER_KEY"):
assert os.environ.get(name), f"{name} is required for the owned OAuth gateway"

8
tests/e2e/migrations/checks.py
View file

@ -3,6 +3,7 @@ from contextlib import ExitStack
from typing import Final
from uuid import uuid4
from e2e_metadata import step
from psycopg import sql
from .containers import Containers, Replica, failed, until
@ -21,12 +22,14 @@ GATED: Final = Migration(
)
@step("start several proxy replica containers")
def start_replicas(
stack: ExitStack, containers: Containers, database: Database, migrations: tuple[Migration, ...] = (), count: int = 3
) -> tuple[Replica, ...]:
return tuple(stack.enter_context(containers.start(database, migrations)) for _ in range(count))
@step("assert the migration completed exactly once")
def assert_completed(database: Database, migration: Migration = COMPLETE) -> None:
assert database.query(
'SELECT finished_at IS NOT NULL, rolled_back_at IS NULL, applied_steps_count FROM '
@ -36,6 +39,7 @@ def assert_completed(database: Database, migration: Migration = COMPLETE) -> Non
assert database.query("SELECT id FROM migration_effect") == ((1,),)
@step("seed a confirmed migration history row")
def confirmed_history(database: Database) -> str:
database.execute(COMPLETE_SQL)
row_id: Final = str(uuid4())
@ -46,6 +50,7 @@ def confirmed_history(database: Database) -> str:
return row_id
@step("assert the original migration proof survived")
def assert_original_proof(database: Database, row_id: str, finished: bool) -> None:
assert database.query(
'SELECT id, applied_steps_count, finished_at IS NOT NULL, rolled_back_at IS NULL FROM '
@ -55,6 +60,7 @@ def assert_original_proof(database: Database, row_id: str, finished: bool) -> No
assert database.query("SELECT id FROM migration_effect") == ((1,),)
@step("pause migration completion with a trigger")
def pause_completion(database: Database) -> None:
database.execute(
sql.SQL(
@ -67,6 +73,7 @@ def pause_completion(database: Database) -> None:
)
@step("crash the migration owner mid-migration")
def interrupt_owner(
containers: Containers, database: Database, after_commit: bool, *, stop_database_session: bool = True
) -> None:
@ -104,6 +111,7 @@ def interrupt_owner(
)
@step("assert replicas refuse an unconfirmed migration")
def unconfirmed(replicas: tuple[Replica, ...], database: Database) -> None:
failed(replicas, "Migration completion could not be verified")
started: Final = str(

12
tests/e2e/migrations/containers.py
View file

@ -11,6 +11,7 @@ from typing import Final
from uuid import uuid4
from e2e_http import NoBody, Success, unwrap
from e2e_metadata import step
from models import KeyGenerateBody, KeyGenerateResponse, KeyInfoParams, KeyInfoResponse
from transport import HttpTransport
@ -20,12 +21,14 @@ from .startup_models import ContainerState, Migration, Observation, Readiness
MASTER_KEY: Final = "sk-migration-ci-fixture"
@step("run a docker command")
def docker(*args: str) -> str:
result: Final = subprocess.run(("docker", *args), capture_output=True, text=True, timeout=90)
assert result.returncode == 0, f"Docker operation failed: {result.stderr}"
return result.stdout.strip()
@step("poll until the condition holds")
def until(description: str, condition: Callable[[], bool], seconds: float = 150) -> None:
deadline: Final = time.monotonic() + seconds
while time.monotonic() < deadline:
@ -41,9 +44,11 @@ class Replica:
transport: HttpTransport
output: Path
@step("inspect the replica container")
def state(self) -> ContainerState:
return ContainerState.model_validate_json(docker("inspect", "--format", "{{json .State}}", self.name))
@step("probe the replica's state and readiness")
def observe(self) -> Observation:
state: Final = self.state()
result: Final = self.transport.get(
@ -52,15 +57,18 @@ class Replica:
ready: Final = isinstance(result, Success) and result.data.status == "healthy" and result.data.db == "connected"
return Observation(None if state.Running else state.ExitCode, ready)
@step("read the replica container logs")
def logs(self) -> str:
result: Final = subprocess.run(("docker", "logs", self.name), capture_output=True, text=True, timeout=30)
assert result.returncode == 0, result.stderr
return result.stdout + result.stderr
@step("kill the replica container")
def kill(self) -> None:
if self.state().Running:
docker("kill", self.name)
@step("confirm the replica mints a usable key")
def usable(self, database: Database) -> None:
alias: Final = f"migration-{uuid4().hex}"
key: Final = unwrap(
@ -86,6 +94,7 @@ class Replica:
) == ((alias,),)
@step("poll replicas for readiness")
def ready(replicas: tuple[Replica, ...], database: Database) -> None:
def all_ready() -> bool:
observations: Final = tuple(replica.observe() for replica in replicas)
@ -97,6 +106,7 @@ def ready(replicas: tuple[Replica, ...], database: Database) -> None:
replica.usable(database)
@step("poll replicas for a rejected startup")
def failed(replicas: tuple[Replica, ...], marker: str) -> None:
def all_stopped() -> bool:
observations: Final = tuple(replica.observe() for replica in replicas)
@ -109,6 +119,7 @@ def failed(replicas: tuple[Replica, ...], marker: str) -> None:
assert marker in replica.logs(), f"Startup failed outside the expected migration: {marker}"
@step("poll replicas to confirm they keep waiting")
def waiting(replicas: tuple[Replica, ...], seconds: float) -> None:
deadline: Final = time.monotonic() + seconds
while time.monotonic() < deadline:
@ -126,6 +137,7 @@ class Containers:
def using(self, image: str) -> "Containers":
return replace(self, image=image)
@step("start a proxy replica container")
@contextmanager
def start(
self,

9
tests/e2e/migrations/database.py
View file

@ -8,6 +8,7 @@ from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from uuid import uuid4
import psycopg
from e2e_metadata import step
from psycopg import sql
from pydantic import TypeAdapter
@ -42,19 +43,23 @@ class Database:
connection.execute("SET statement_timeout = '15s'")
yield connection
@step("run a SQL statement")
def execute(self, statement: LiteralString | sql.Composed, params: tuple[Scalar, ...] = ()) -> None:
with self.connection() as connection:
connection.execute(statement, params or None)
@step("query the database")
def query(
self, statement: LiteralString | sql.Composed, params: tuple[Scalar, ...] = ()
) -> tuple[tuple[Scalar, ...], ...]:
with self.connection() as connection:
return ROWS.validate_python(connection.execute(statement, params or None).fetchall())
@step("check whether a table exists")
def exists(self, name: str) -> bool:
return self.query("SELECT to_regclass(%s) IS NOT NULL", (name,)) == ((True,),)
@step("read the _prisma_migrations history")
def history(self) -> tuple[tuple[Scalar, ...], ...]:
if not self.exists("_prisma_migrations"):
return ()
@ -63,6 +68,7 @@ class Database:
"applied_steps_count, logs FROM _prisma_migrations ORDER BY id"
)
@step("list backends waiting on an advisory lock")
def blocked(self, key: int = GATE_KEY) -> tuple[tuple[Scalar, ...], ...]:
return self.query(
"SELECT pid FROM pg_locks WHERE locktype = 'advisory' AND NOT granted "
@ -71,6 +77,7 @@ class Database:
(key >> 32, key & 0xFFFFFFFF),
)
@step("hold an advisory lock")
@contextmanager
def lock(self, key: int = GATE_KEY) -> Generator[None]:
with self.connection() as connection:
@ -86,6 +93,7 @@ class Databases:
admin_url: str
container_admin_url: str
@step("create a test database")
@contextmanager
def create(self, template: Database | None = None, schema: str = "public") -> Generator[Database]:
name: Final = f"litellm_migration_test_{uuid4().hex[:20]}"
@ -105,6 +113,7 @@ class Databases:
connection.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name)))
@step("create a read-only database role")
@contextmanager
def restricted_user(database: Database) -> Generator[Database]:
role: Final = f"migration_reader_{uuid4().hex[:16]}"

7
tests/e2e/migrations/upgrade.py
View file

@ -8,6 +8,7 @@ from typing import Final
from uuid import uuid4
from e2e_http import Result, Success, unwrap
from e2e_metadata import step
from models import (
KeyGenerateBody,
KeyGenerateResponse,
@ -24,6 +25,7 @@ from .database import Database
CACHED_PLAN: Final = "cached plan must not change result type"
@step("generate a virtual key on the replica")
def provision(replica: Replica) -> tuple[str, str]:
alias: Final = f"upgrade-{uuid4().hex}"
key: Final = unwrap(
@ -37,6 +39,7 @@ def provision(replica: Replica) -> tuple[str, str]:
return key, alias
@step("confirm the key resolves on the replica")
def confirm(replica: Replica, key: str, alias: str) -> None:
info: Final = unwrap(
replica.transport.get(
@ -62,6 +65,7 @@ class Outcomes:
self.failures.append(result.model_dump_json())
@step("drive virtual-key auth traffic")
@contextmanager
def auth_traffic(replica: Replica, key: str, interval: float = 0.05) -> Generator[Outcomes]:
outcomes: Final = Outcomes()
@ -93,6 +97,7 @@ def auth_traffic(replica: Replica, key: str, interval: float = 0.05) -> Generato
)
@step("poll auth traffic for more served requests")
def keep_serving(outcomes: Outcomes, description: str, calls: int = 20) -> int:
target: Final = outcomes.served + calls
until(description, lambda: outcomes.served >= target or bool(outcomes.failures))
@ -100,10 +105,12 @@ def keep_serving(outcomes: Outcomes, description: str, calls: int = 20) -> int:
return outcomes.served
@step("read applied migration names")
def migration_names(database: Database) -> frozenset[str]:
return frozenset(str(row[0]) for row in database.query("SELECT migration_name FROM _prisma_migrations"))
@step("assert migration history is clean")
def assert_history_clean(database: Database) -> None:
assert database.query(
"SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS NOT NULL"

3
tests/e2e/otel_client.py
View file

@ -25,6 +25,7 @@ from pydantic import BaseModel, ConfigDict, Field
from e2e_config import OTEL_QUERY_URL, POLL_INTERVAL, POLL_TIMEOUT
from e2e_http import URL, NetworkError, NoBody, Result, Success, get
from e2e_metadata import step
#: OTEL resource service.name the proxy exports under (OTEL_SERVICE_NAME default).
JAEGER_SERVICE = "litellm"
@ -109,6 +110,7 @@ class OtelReader:
timeout=30.0,
)
@step("query Jaeger for the call's traces")
def traces_for_call(self, call_id: str) -> list[JaegerTrace]:
"""Every trace holding a span tagged with this call id. Jaeger matches
spans server-side and returns their full traces; more than one hit for
@ -119,6 +121,7 @@ class OtelReader:
case failure:
pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}")
@step("poll Jaeger for the call's settled trace")
def poll_traces_for_call(
self, *, call_id: str, settled_names: set[str], settled_prefixes: set[str]
) -> list[JaegerTrace]:

6
tests/e2e/other/other_client.py
View file

@ -16,6 +16,7 @@ from __future__ import annotations
from dataclasses import dataclass
from e2e_http import NoBody, ProbeResult, Result
from e2e_metadata import step
from idp import Keycloak, keycloak_from_env
from models import (
ReadinessDetailsResponse,
@ -35,11 +36,13 @@ class OtherClient:
"""Resolved per use, so the suite's non-JWT tests never need the IdP env."""
return keycloak_from_env()
@step("GET /health/liveliness")
def liveness(self) -> ProbeResult:
"""GET /health/liveliness. Unauthenticated; the probe returns status +
raw body so the test can assert the worker reports itself alive."""
return self.proxy.transport.probe("/health/liveliness", params=NoBody())
@step("GET /health/readiness without auth")
def readiness_public(self) -> Result[ReadinessResponse]:
"""GET /health/readiness with no credential at all, proving the probe is
safe to expose to an unauthenticated load balancer."""
@ -50,6 +53,7 @@ class OtherClient:
response_type=ReadinessResponse,
)
@step("GET /health/readiness/details")
def readiness_details(self, key: str) -> Result[ReadinessDetailsResponse]:
return self.proxy.transport.get(
"/health/readiness/details",
@ -58,6 +62,7 @@ class OtherClient:
response_type=ReadinessDetailsResponse,
)
@step("GET /health/readiness/details without auth")
def readiness_details_unauthenticated(self) -> Result[ReadinessDetailsResponse]:
return self.proxy.transport.get(
"/health/readiness/details",
@ -66,6 +71,7 @@ class OtherClient:
response_type=ReadinessDetailsResponse,
)
@step("GET /user/list")
def list_users_as(self, key: str) -> Result[UserListResponse]:
"""GET /user/list under `key`. Admin-only, so it doubles as the master
key's authorization proof: the master key (proxy admin) reads it, a

49
tests/e2e/proxy_client.py
View file

@ -42,6 +42,7 @@ from e2e_http import (
is_ok,
unwrap,
)
from e2e_metadata import STEP_FRAMES, step
from models import (
AnthropicMessagesBody,
AnthropicMessagesResponse,
@ -467,6 +468,7 @@ class ProxyClient:
# ---- keys / customers (satisfies lifecycle.ResourceClient) ----------
@step("generate virtual key")
def generate_key(self, body: KeyGenerateBody) -> str:
return unwrap(
self.transport.post(
@ -477,6 +479,7 @@ class ProxyClient:
)
).key
@step("delete virtual key")
def delete_key(self, key: str) -> None:
_ = self.transport.post(
"/key/delete",
@ -485,6 +488,7 @@ class ProxyClient:
response_type=NoBody,
)
@step("delete end users")
def delete_customers(self, user_ids: list[str]) -> None:
if not user_ids:
return
@ -495,6 +499,7 @@ class ProxyClient:
response_type=NoBody,
)
@step("read /key/info")
def key_info(self, key: str) -> KeyInfo:
return unwrap(
self.transport.get(
@ -505,6 +510,7 @@ class ProxyClient:
)
).info
@step("read /debug/memory/summary on every replica")
def memory_summary_everywhere(self) -> Mapping[str, Result[MemorySummaryResponse]]:
return {
url: transport.get(
@ -516,6 +522,7 @@ class ProxyClient:
for url, transport in self.replicas.items()
}
@step("poll every replica for the converged read")
def read_back_everywhere[R: BaseModel](
self,
path: str,
@ -563,6 +570,7 @@ class ProxyClient:
path, headers=self.management_headers(transport=transport), params=params, response_type=response_type
)
@step("read /model/info")
def model_info(self) -> list[ModelInfoEntry]:
"""Every configured deployment with the price the proxy resolved for it
(config override merged over cost-map defaults)."""
@ -575,6 +583,7 @@ class ProxyClient:
)
).data
@step("read /router/settings")
def router_settings(self) -> RouterCurrentValues:
"""The router knobs the proxy is running with, for a test whose behavior
needs one of them switched on in the proxy config."""
@ -587,6 +596,7 @@ class ProxyClient:
)
).current_values
@step("read the public model cost map")
def model_cost_map(self) -> dict[str, CostMapEntry]:
return unwrap(
self.transport.get(
@ -597,6 +607,7 @@ class ProxyClient:
)
).root
@step("GET /v1/files")
def list_files(self, key: str) -> Result[FileListResponse]:
return self.transport.get(
"/v1/files",
@ -605,6 +616,7 @@ class ProxyClient:
response_type=FileListResponse,
)
@step("GET /v1/fine_tuning/jobs")
def list_fine_tuning_jobs(self, key: str, params: FineTuningJobsParams) -> Result[FineTuningJobsResponse]:
return self.transport.get(
"/v1/fine_tuning/jobs",
@ -613,6 +625,7 @@ class ProxyClient:
response_type=FineTuningJobsResponse,
)
@step("register deployment")
def create_model(
self,
model_name: str,
@ -632,6 +645,7 @@ class ProxyClient:
provider_live=provider_live,
)
@step("check a general_settings flag")
def general_setting_enabled(self, field_name: str) -> bool:
"""Whether the proxy is running with the named general_settings flag on, for
a test whose behavior only exists under a config flag the stack has to carry."""
@ -645,6 +659,7 @@ class ProxyClient:
).root
return any(entry.field_name == field_name and entry.field_value is True for entry in fields)
@step("register deployment")
def register_model(
self, body: ModelNewBody, listed_for: str | None = None, *, provider_live: bool = False
) -> str:
@ -727,6 +742,7 @@ class ProxyClient:
timeout=poll_timeout,
)
@step("update deployment")
def update_model(self, model_id: str, litellm_params: LiteLLMParamsBody) -> None:
"""Merge `litellm_params` over the deployment `model_id`'s stored params via
POST /model/update. The proxy overlays only the non-null fields and clears
@ -744,6 +760,7 @@ class ProxyClient:
)
)
@step("delete deployment")
def delete_model(self, model_id: str) -> None:
result = self.transport.post(
"/model/delete",
@ -752,7 +769,7 @@ class ProxyClient:
response_type=NoBody,
)
if not is_ok(result):
warnings.warn(f"delete_model({model_id!r}) failed: {result}", stacklevel=2)
warnings.warn(f"delete_model({model_id!r}) failed: {result}", stacklevel=2 + STEP_FRAMES)
# ---- replica read-back ----------------------------------------------
@ -770,6 +787,7 @@ class ProxyClient:
assert replicas, f"no replica is configured to serve {path}, so a read-back there would prove nothing"
return replicas
@step("poll every replica for the settled body")
def read_body_back_everywhere[R: BaseModel](
self, path: str, response_type: type[R], *, settled: Callable[[R], bool]
) -> Mapping[str, R]:
@ -795,6 +813,7 @@ class ProxyClient:
f"last read: {last}"
)
@step("poll every replica for a 404")
def gone_everywhere(self, path: str) -> Mapping[str, int]:
"""Poll GET `path` on every replica that serves it until each stops serving
it, and fail naming the first replica that still does at poll_timeout.
@ -827,6 +846,7 @@ class ProxyClient:
# ---- mcp toolsets ---------------------------------------------------
@step("create MCP toolset")
def create_toolset(self, body: ToolsetCreateBody) -> ToolsetRow:
return unwrap(
self.transport.post(
@ -837,6 +857,7 @@ class ProxyClient:
)
)
@step("update MCP toolset")
def update_toolset(self, body: ToolsetUpdateBody) -> ToolsetRow:
"""PUT /v1/mcp/toolset: a partial update where a field left unset keeps its
stored value and None clears it."""
@ -849,6 +870,7 @@ class ProxyClient:
)
)
@step("delete MCP toolset")
def delete_toolset(self, toolset_id: str) -> Result[NoBody]:
"""DELETE /v1/mcp/toolset/{toolset_id}. Returns the outcome so the act phase
can unwrap it while a deferred teardown can ignore an already-deleted row."""
@ -859,6 +881,7 @@ class ProxyClient:
response_type=NoBody,
)
@step("create credential")
def create_credential(self, body: CredentialCreateBody) -> None:
unwrap(
self.transport.post(
@ -869,6 +892,7 @@ class ProxyClient:
)
)
@step("delete credential")
def delete_credential(self, credential_name: str) -> None:
result = self.transport.delete(
f"/credentials/{credential_name}",
@ -877,8 +901,9 @@ class ProxyClient:
response_type=NoBody,
)
if not is_ok(result):
warnings.warn(f"delete_credential({credential_name!r}) failed: {result}", stacklevel=2)
warnings.warn(f"delete_credential({credential_name!r}) failed: {result}", stacklevel=2 + STEP_FRAMES)
@step("create team")
def create_team(self, body: TeamNewBody) -> str:
return unwrap(
self.transport.post(
@ -889,6 +914,7 @@ class ProxyClient:
)
).team_id
@step("update team")
def update_team(self, body: TeamUpdateBody) -> None:
unwrap(
self.transport.post(
@ -899,6 +925,7 @@ class ProxyClient:
)
)
@step("delete team")
def delete_team(self, team_id: str) -> None:
result = self.transport.post(
"/team/delete",
@ -907,8 +934,9 @@ class ProxyClient:
response_type=NoBody,
)
if not is_ok(result):
warnings.warn(f"delete_team({team_id!r}) failed: {result}", stacklevel=2)
warnings.warn(f"delete_team({team_id!r}) failed: {result}", stacklevel=2 + STEP_FRAMES)
@step("delete internal user")
def delete_user(self, user_id: str) -> None:
"""Best-effort teardown; a 404 is not a leak, since JWT tests defer this for
a user the proxy only upserts after a successful auth."""
@ -922,10 +950,11 @@ class ProxyClient:
case Success() | UnknownApiError(status_code=404):
return
case _:
warnings.warn(f"delete_user({user_id!r}) failed: {result}", stacklevel=2)
warnings.warn(f"delete_user({user_id!r}) failed: {result}", stacklevel=2 + STEP_FRAMES)
# ---- LLM calls ------------------------------------------------------
@step("POST /chat/completions")
def chat(self, key: str, body: ChatBody) -> Result[ChatResponse]:
return self.transport.post(
"/chat/completions",
@ -934,12 +963,15 @@ class ProxyClient:
response_type=ChatResponse,
)
@step("POST /chat/completions (streaming)")
def chat_stream(self, key: str, body: ChatBody) -> StreamingResponse:
return self.transport.stream("/chat/completions", headers=self.transport.bearer(key), json=body)
@step("POST /v1/messages (streaming)")
def messages_stream(self, key: str, body: AnthropicMessagesBody) -> StreamingResponse:
return self.transport.stream("/v1/messages", headers=self.transport.bearer(key), json=body)
@step("POST /embeddings")
def embed(self, key: str, body: EmbedBody) -> Result[EmbedResponse]:
return self.transport.post(
"/embeddings",
@ -948,6 +980,7 @@ class ProxyClient:
response_type=EmbedResponse,
)
@step("POST /v1/ocr")
def ocr(self, key: str, body: OcrBody) -> Result[OcrResponse]:
return self.transport.post(
"/v1/ocr",
@ -957,6 +990,7 @@ class ProxyClient:
timeout=SLOW_PROVIDER_TIMEOUT_SECONDS,
)
@step("POST /v1/rerank")
def rerank(self, key: str, body: RerankBody) -> Result[RerankResponse]:
"""POST /v1/rerank (Cohere-format). No official OpenAI/Anthropic SDK
covers this route, so it stays on the shared typed transport."""
@ -967,6 +1001,7 @@ class ProxyClient:
response_type=RerankResponse,
)
@step("POST /v1/messages/count_tokens")
def count_tokens(self, key: str, body: CountTokensBody) -> Result[CountTokensResponse]:
"""POST /v1/messages/count_tokens (Anthropic-native). Sends the
anthropic-version header so the native path accepts it; harmless on the
@ -978,6 +1013,7 @@ class ProxyClient:
response_type=CountTokensResponse,
)
@step("POST /v1/messages")
def messages(self, key: str, body: AnthropicMessagesBody) -> Result[AnthropicMessagesResponse]:
"""POST /v1/messages (Anthropic-native). The response is either the
Anthropic-shape passthrough (`content`) or the OpenAI-normalized shape
@ -994,6 +1030,7 @@ class ProxyClient:
# ---- spend read-back ------------------------------------------------
@step("GET /spend/logs")
def spend_logs(self, params: SpendLogsParams) -> list[SpendLogRow]:
result = self.transport.get(
"/spend/logs",
@ -1007,6 +1044,7 @@ class ProxyClient:
case _:
return []
@step("read spend logs for a time window")
def spend_logs_window(self, *, start: datetime, end: datetime) -> list[SpendLogRow]:
def fetch(page: int) -> SpendLogsPage:
return unwrap(
@ -1029,11 +1067,13 @@ class ProxyClient:
*(row for page in range(2, first.total_pages + 1) for row in fetch(page).data),
]
@step("poll /spend/logs for the key")
def poll_logs_for_key(
self, key: str, *, min_rows: int = 1, predicate: RowsPredicate | None = None
) -> list[SpendLogRow]:
return self._poll(lambda: self.spend_logs(SpendLogsParams(api_key=key)), min_rows, predicate)
@step("poll /spend/logs for the request id")
def poll_logs_for_request_id(
self,
request_id: str,
@ -1064,6 +1104,7 @@ class ProxyClient:
# ---- route probe ----------------------------------------------------
@step("probe a management route")
def probe(self, path: str, *, params: NoBody) -> ProbeResult:
return self.transport.probe(path, params=params, headers=self.management_headers())

20
tests/e2e/quota_management/budgets/budget_client.py
View file

@ -17,6 +17,7 @@ from datetime import datetime
from pydantic import AliasPath, BaseModel, Field, RootModel
from e2e_http import NoBody, Result, StreamingResponse, Success, unwrap
from e2e_metadata import step
from proxy_client import ProxyClient
from models import (
AnthropicMessagesBody,
@ -224,6 +225,7 @@ class BudgetClient:
# ---- generic key ops (delegate to the shared ProxyClient) ---------------
@step("generate virtual key")
def generate_key(
self,
*,
@ -256,12 +258,14 @@ class BudgetClient:
def delete_key(self, key: str) -> None:
self.proxy.delete_key(key)
@step("read the key's budget windows")
def key_budget_windows(self, key: str) -> list[BudgetWindowState]:
"""A key's budget_limits windows as /key/info stores them. Each window's
reset_at is advanced by the reset job in the same pass that zeroes the
window's spend counter, so a strictly-later value proves the wipe ran."""
return self.proxy.key_info(key).budget_limits or []
@step("read the team's budget windows")
def team_budget_windows(self, team_id: str) -> list[BudgetWindowState]:
"""Team analog of key_budget_windows, read from /team/info."""
match self._team_info(team_id):
@ -275,6 +279,7 @@ class BudgetClient:
# ---- chat (raw HTTP outcome: a budget block surfaces as a non-2xx) --
@step("POST /chat/completions")
def chat(
self,
key: str,
@ -297,6 +302,7 @@ class BudgetClient:
),
)
@step("POST /v1/messages")
def messages(
self,
key: str,
@ -317,6 +323,7 @@ class BudgetClient:
# ---- internal user --------------------------------------------------
@step("create internal user with a budget")
def create_user(self, *, max_budget: float, budget_duration: str | None = None) -> str:
return unwrap(
self.proxy.transport.post(
@ -335,6 +342,7 @@ class BudgetClient:
response_type=NoBody,
)
@step("read /user/info for recorded spend")
def user_info(self, user_id: str) -> UserInfoRow | None:
result = self.proxy.transport.get(
"/user/info",
@ -350,6 +358,7 @@ class BudgetClient:
# ---- customer / end-user -------------------------------------------
@step("create end user with a budget")
def create_customer(
self,
customer_id: str,
@ -369,6 +378,7 @@ class BudgetClient:
# ---- organization ---------------------------------------------------
@step("create organization with a budget")
def create_org(self, *, max_budget: float, alias: str, budget_duration: str | None = None) -> str:
return unwrap(
self.proxy.transport.post(
@ -383,6 +393,7 @@ class BudgetClient:
)
).organization_id
@step("read the organization's budget id")
def org_budget_id(self, org_id: str) -> str | None:
"""The id of the budget row backing an org; its budget_reset_at is read via
budget_info (LIT-4570: /organization/new stores budget_duration without
@ -409,6 +420,7 @@ class BudgetClient:
# ---- team -----------------------------------------------------------
@step("create team with a budget")
def create_team(
self,
*,
@ -463,6 +475,7 @@ class BudgetClient:
assert last is not None
raise AssertionError(last)
@step("add a member to the team")
def add_team_member(self, team_id: str, user_id: str, *, max_budget_in_team: float | None = None) -> None:
last_body = ""
for attempt in range(_TEAM_READY_ATTEMPTS):
@ -484,6 +497,7 @@ class BudgetClient:
break
raise AssertionError(last_body)
@step("update the member's in-team budget")
def update_team_member(
self,
team_id: str,
@ -504,6 +518,7 @@ class BudgetClient:
)
assert resp.ok, resp.body
@step("read the member budget's reset_at")
def member_budget_reset_at(self, team_id: str, user_id: str) -> str | None:
"""The member's per-team budget_reset_at as /team/info reports it, or None if
no reset is scheduled. The reset job advances this each time the window
@ -519,6 +534,7 @@ class BudgetClient:
# ---- tag ------------------------------------------------------------
@step("create tag with a budget")
def create_tag(self, name: str, *, max_budget: float) -> str:
resp = self.proxy.transport.send(
"/tag/new",
@ -538,6 +554,7 @@ class BudgetClient:
# ---- model access group ---------------------------------------------
@step("set the model access group's budget")
def set_access_group_budget(
self,
access_group: str,
@ -561,6 +578,7 @@ class BudgetClient:
)
)
@step("read the model access group's budget")
def access_group_budget(self, access_group: str) -> AccessGroupBudgetResponse:
return unwrap(
self.proxy.transport.get(
@ -581,6 +599,7 @@ class BudgetClient:
# ---- budget table ---------------------------------------------------
@step("create a shared budget")
def create_budget(
self,
*,
@ -611,6 +630,7 @@ class BudgetClient:
response_type=NoBody,
)
@step("read /budget/info")
def budget_info(self, budget_id: str) -> tuple[BudgetRow, ...]:
result = self.proxy.transport.post(
"/budget/info",

2
tests/e2e/quota_management/ratelimit/quota_client.py
View file

@ -9,6 +9,7 @@ from dataclasses import dataclass
from proxy_client import ProxyClient
from e2e_http import StreamingResponse
from e2e_metadata import step
from models import ChatBody, ChatMessage
@ -16,6 +17,7 @@ from models import ChatBody, ChatMessage
class QuotaClient:
proxy: ProxyClient
@step("POST /chat/completions")
def chat(self, key: str, model: str, content: str, *, max_tokens: int = 16) -> StreamingResponse:
return self.proxy.transport.send(
"/chat/completions",

6
tests/e2e/quota_management/spend_tracking/cost_rows.py
View file

@ -36,6 +36,7 @@ from pydantic import BaseModel, RootModel
from e2e_config import unique_marker
from e2e_http import Success
from e2e_metadata import step
from lifecycle import ResourceManager
from models import LiteLLMParamsBody, SpendLogsParams
from proxy_client import ProxyClient
@ -98,6 +99,7 @@ def approx_equal(actual: float, expected: float) -> bool:
return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)
@step("assert the row's total is the sum of its components")
def assert_total_is_sum_of_components(row: CostRow) -> None:
"""The row's total is input + output + tool usage. The cache components are
already inside the gross input cost, so adding them again would double-bill."""
@ -114,6 +116,7 @@ def assert_total_is_sum_of_components(row: CostRow) -> None:
)
@step("assert fresh input tokens are billed at the model's rate")
def assert_fresh_tokens_billed_at(row: CostRow, input_rate: float) -> None:
"""Strip the cache components out of the gross input cost and what is left must
be the freshly-read tokens at the deployment's input rate."""
@ -133,6 +136,7 @@ def assert_fresh_tokens_billed_at(row: CostRow, input_rate: float) -> None:
)
@step("poll /spend/logs for the request's cost row")
def poll_cost_row(proxy: ProxyClient, request_id: str) -> CostRow | None:
"""Poll /spend/logs for the call's row until it lands with a cost breakdown
(rows flush ~60s behind the call via proxy_batch_write_at); None on timeout."""
@ -156,6 +160,7 @@ def poll_cost_row(proxy: ProxyClient, request_id: str) -> CostRow | None:
return None
@step("poll /spend/logs for a matching cost row")
def poll_cost_row_where(
proxy: ProxyClient, api_key: str, predicate: Callable[[CostRow], bool]
) -> CostRow | None:
@ -182,6 +187,7 @@ def poll_cost_row_where(
return None
@step("register a deployment with custom pricing")
def register_priced_model(
proxy: ProxyClient,
resources: ResourceManager,

28
tests/e2e/quota_management/spend_tracking/spend_e2e_client.py
View file

@ -29,6 +29,7 @@ from e2e_http import (
is_ok,
unwrap,
)
from e2e_metadata import step
from models import (
AnthropicMessagesBody,
ChatBody,
@ -234,6 +235,7 @@ def _chat_body(
class SpendClient:
proxy: ProxyClient
@step("POST /chat/completions")
def chat(
self,
key: str,
@ -250,6 +252,7 @@ class SpendClient:
_chat_body(model, content, max_tokens=max_tokens, tags=tags, user=user, cache=cache),
)
@step("POST /chat/completions (streaming)")
def chat_stream(
self, key: str, model: str, content: str, *, max_tokens: int | None = None
) -> StreamingResponse:
@ -257,6 +260,7 @@ class SpendClient:
key, _chat_body(model, content, max_tokens=max_tokens, stream=True)
)
@step("POST /v1/messages (streaming)")
def messages_stream(
self, key: str, model: str, content: str, *, max_tokens: int
) -> StreamingResponse:
@ -270,9 +274,11 @@ class SpendClient:
),
)
@step("POST /embeddings")
def embed(self, key: str, model: str, content: str) -> Result[EmbedResponse]:
return self.proxy.embed(key, EmbedBody(model=model, input=content))
@step("poll /spend/logs for the key")
def poll_logs_for_key(
self,
key: str,
@ -284,6 +290,7 @@ class SpendClient:
key, min_rows=min_rows, predicate=predicate
)
@step("POST /spend/calculate")
def calculate_spend(self, model: str, content: str) -> float:
return unwrap(
self.proxy.transport.post(
@ -296,6 +303,7 @@ class SpendClient:
)
).cost
@step("GET /spend/tags")
def spend_by_tags(self) -> list[TagSpend]:
result = self.proxy.transport.get(
"/spend/tags",
@ -309,6 +317,7 @@ class SpendClient:
case _:
return []
@step("poll /spend/tags for the tag")
def poll_tag_spend(self, tag: str, *, minimum: float = 0.0) -> TagSpend | None:
"""Poll /spend/tags until the tag's aggregate reaches `minimum`; last seen."""
deadline = time.monotonic() + self.proxy.poll_timeout
@ -324,6 +333,7 @@ class SpendClient:
time.sleep(self.proxy.poll_interval)
return entry
@step("poll /key/info for recorded spend")
def poll_key_spend(self, key: str, *, minimum: float = 0.0) -> float:
deadline = time.monotonic() + self.proxy.poll_timeout
spend = 0.0
@ -334,6 +344,7 @@ class SpendClient:
time.sleep(self.proxy.poll_interval)
return spend
@step("GET /spend/logs")
def spend_logs_page(
self, *, api_key: str | None, page: int, page_size: int
) -> SpendLogsPage:
@ -356,6 +367,7 @@ class SpendClient:
)
)
@step("probe the spend route")
def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult:
return self.proxy.transport.probe(path, params=params)
@ -370,6 +382,7 @@ class SpendClient:
)
return outcome.result if isinstance(outcome, Converged) else outcome.last_result
@step("create internal user")
def create_user(self, *, email: str, role: UserRole, user_id: str) -> str:
return unwrap(
self.proxy.transport.post(
@ -390,6 +403,7 @@ class SpendClient:
)
)
@step("generate virtual key")
def generate_key_record(self, body: KeyGenerateBody) -> KeyGenerateResponse:
return unwrap(
self.proxy.transport.post(
@ -400,6 +414,7 @@ class SpendClient:
)
)
@step("POST /chat/completions")
def send_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
return self.proxy.transport.send(
"/chat/completions",
@ -407,6 +422,7 @@ class SpendClient:
json=_chat_body(model, content, max_tokens=max_tokens),
)
@step("POST /chat/completions (queued)")
def send_queued_chat(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
return self.proxy.transport.send(
"/queue/chat/completions",
@ -418,6 +434,7 @@ class SpendClient:
),
)
@step("POST /v1/messages")
def send_messages(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
return self.proxy.transport.send(
"/v1/messages",
@ -429,6 +446,7 @@ class SpendClient:
),
)
@step("POST /v1/responses")
def send_responses(self, key: str, model: str, content: str) -> StreamingResponse:
return self.proxy.transport.send(
"/v1/responses",
@ -436,6 +454,7 @@ class SpendClient:
json=ResponsesBody(model=model, input=content),
)
@step("POST /embeddings")
def send_embed(self, key: str, model: str, content: str) -> StreamingResponse:
return self.proxy.transport.send(
"/embeddings",
@ -443,6 +462,7 @@ class SpendClient:
json=EmbedBody(model=model, input=content),
)
@step("POST the gemini passthrough generateContent")
def send_gemini_generate(self, key: str, model: str, content: str, *, max_tokens: int) -> StreamingResponse:
return self.proxy.transport.send(
f"/gemini/v1beta/models/{model}:generateContent",
@ -453,6 +473,7 @@ class SpendClient:
),
)
@step("POST /v1/files")
def upload_batch_file(self, key: str, model: str, content: bytes) -> FileObject:
return unwrap(
self.proxy.transport.upload(
@ -466,6 +487,7 @@ class SpendClient:
)
)
@step("POST /v1/batches")
def create_batch(self, key: str, body: BatchCreateBody) -> BatchObject:
return unwrap(
self.proxy.transport.post(
@ -476,6 +498,7 @@ class SpendClient:
)
)
@step("GET /v1/batches/{id}")
def retrieve_batch(self, key: str, batch_id: str, *, provider: str) -> BatchObject:
return unwrap(
self.proxy.transport.get(
@ -486,6 +509,7 @@ class SpendClient:
)
)
@step("replay a provider callback log")
def replay_callback_log(self, key: str, payload: CallbackLogPayload) -> CallbackLogsResponse:
return unwrap(
self.proxy.transport.post(
@ -496,9 +520,11 @@ class SpendClient:
)
)
@step("GET /health for the deployment")
def health(self, model: str) -> ProbeResult:
return self.proxy.transport.probe("/health", params=HealthParams(model=model))
@step("GET /user/daily/activity")
def daily_activity_for_key(self, token: str, *, start: datetime, end: datetime) -> DailyActivityKeyBreakdown | None:
response: Final = unwrap(
self.proxy.transport.get(
@ -517,6 +543,7 @@ class SpendClient:
None,
)
@step("poll /user/daily/activity for the key")
def poll_daily_activity_for_key(
self, token: str, *, start: datetime, end: datetime, min_requests: int
) -> DailyActivityKeyBreakdown | None:
@ -530,6 +557,7 @@ class SpendClient:
)
return outcome.result if isinstance(outcome, Converged) else outcome.last_result
@step("GET the proxy's OpenAPI schema")
def openapi(self) -> OpenAPISchema:
return unwrap(
self.proxy.transport.get(

3
tests/e2e/quota_management/spend_tracking/spend_reconciliation.py
View file

@ -5,6 +5,7 @@ from typing import Final
from e2e_config import provider_edge_base, unique_marker
from e2e_http import unwrap
from e2e_metadata import step
from lifecycle import ResourceManager
from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody, LiteLLMParamsBody, TeamNewBody
from spend_e2e_client import SpendClient
@ -32,6 +33,7 @@ class TeamTraffic:
return self.prompt_tokens * INPUT_RATE + self.completion_tokens * OUTPUT_RATE
@step("send chat traffic from two teams")
def create_traffic(client: SpendClient, resources: ResourceManager) -> tuple[TeamTraffic, ...]:
base: Final = provider_edge_base("openai")
model: Final = f"e2e-reconciliation-{unique_marker()}"
@ -84,6 +86,7 @@ def create_traffic(client: SpendClient, resources: ResourceManager) -> tuple[Tea
return tuple(team_traffic() for _ in range(2))
@step("reconcile /spend/logs against the responses")
def assert_logs_match(client: SpendClient, traffic: TeamTraffic) -> None:
expected_ids: Final = frozenset(response.id for response in traffic.responses)
assert len(expected_ids) == len(traffic.responses), "responses must have distinct IDs"

18
tests/e2e/router/reliability_support.py
View file

@ -22,6 +22,7 @@ from pydantic import ValidationError
from proxy_client import ProxyClient
from e2e_config import CHEAP_OPENAI_MODEL, PROXY_BASE_URL, unique_marker
from e2e_http import NetworkError, StreamHead, StreamingResponse
from e2e_metadata import step
from models import (
CacheControl,
ChatMessage,
@ -78,6 +79,7 @@ def cached_system_turn(marker: str) -> ChatMessage:
return ChatMessage(role="system", content=[TextContentPart(text=filler, cache_control=CacheControl())])
@step("register deployment with an unreachable base")
def create_bad_base_deployment(proxy: ProxyClient, name: str) -> str:
"""Register a deployment pointing at an unreachable base, so every call to it
fails with a real connection error the fallback can reroute around."""
@ -86,6 +88,7 @@ def create_bad_base_deployment(proxy: ProxyClient, name: str) -> str:
)
@step("register never-benched unreachable deployment")
def create_never_benched_refusing_deployment(proxy: ProxyClient, name: str) -> str:
return proxy.create_model(
name,
@ -93,17 +96,20 @@ def create_never_benched_refusing_deployment(proxy: ProxyClient, name: str) -> s
)
@step("register deployment with a 1ms timeout")
def create_timeout_deployment(proxy: ProxyClient, name: str) -> str:
"""Register a deployment with a 1ms deadline the real backend always exceeds."""
return proxy.create_model(name, LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001))
@step("register small-context deployment")
def create_small_context_deployment(proxy: ProxyClient, name: str) -> str:
"""Register a deployment on the smallest-context model OpenAI still serves, so an
oversized prompt earns a real context-window refusal from the provider."""
return proxy.create_model(name, LiteLLMParamsBody(model=SMALL_CONTEXT_MODEL, api_key=REAL_KEY))
@step("register content-filtered Azure deployment")
def create_content_filtered_deployment(proxy: ProxyClient, name: str) -> str:
"""Register the Azure OpenAI deployment whose content filter refuses
CONTENT_POLICY_PROMPT with a real policy-violation 400 (the one live trigger
@ -121,6 +127,7 @@ def create_content_filtered_deployment(proxy: ProxyClient, name: str) -> str:
)
@step("register Azure deployment with zero allowed fails")
def create_azure_benched_on_first_failure_deployment(proxy: ProxyClient, name: str, cooldown_time: float) -> str:
"""The live Azure OpenAI deployment holding all of the group's shuffle weight,
benched on its first failure of any class, with the client's own retries off."""
@ -141,6 +148,7 @@ def create_azure_benched_on_first_failure_deployment(proxy: ProxyClient, name: s
)
@step("register prompt-caching Anthropic deployment")
def create_caching_deployment(proxy: ProxyClient, name: str) -> str:
"""Register the Anthropic deployment whose prompt cache the affinity check pins to."""
return proxy.create_model(name, LiteLLMParamsBody(model=CACHING_MODEL, api_key=CACHING_KEY, weight=1))
@ -161,6 +169,7 @@ def _register_benched_on_first_failure(
)
@step("register always-timing-out deployment")
def create_always_timing_out_deployment(proxy: ProxyClient, name: str, cooldown_time: float | None = None) -> str:
"""A 1ms deadline the real backend always exceeds, benched on its first Timeout."""
return _register_benched_on_first_failure(
@ -171,6 +180,7 @@ def create_always_timing_out_deployment(proxy: ProxyClient, name: str, cooldown_
)
@step("register always-401 deployment")
def create_always_unauthorized_deployment(proxy: ProxyClient, name: str, cooldown_time: float | None = None) -> str:
"""A key the real backend rejects with a 401, benched on its first AuthenticationError."""
return _register_benched_on_first_failure(
@ -199,6 +209,7 @@ def _nested_proxy_params(upstream_group: str, upstream_key: str, cooldown_time:
)
@step("register always-5xx deployment")
def create_always_5xx_deployment(
proxy: ProxyClient, name: str, upstream_group: str, upstream_key: str, cooldown_time: float | None = None
) -> str:
@ -212,6 +223,7 @@ def create_always_5xx_deployment(
)
@step("register always-429 deployment")
def create_always_rate_limited_deployment(
proxy: ProxyClient, name: str, upstream_group: str, upstream_key: str, cooldown_time: float | None = None
) -> str:
@ -222,6 +234,7 @@ def create_always_rate_limited_deployment(
)
@step("use up the rpm-limited key's one request")
def spend_only_request_of(proxy: ProxyClient, spent_key: str) -> None:
"""Uses up the one request an rpm_limit=1 key allows. The proxy's rate limiter
opens the key's 60s window on this call, so it goes right before the calls that
@ -234,6 +247,7 @@ def spend_only_request_of(proxy: ProxyClient, spent_key: str) -> None:
)
@step("register always-picked small-context deployment")
def create_always_picked_small_context_deployment(proxy: ProxyClient, name: str) -> str:
"""The always-picked half of a retry pair on the smallest-context model OpenAI
still serves: it holds all of the model group's shuffle weight, so an oversized
@ -248,6 +262,7 @@ def create_always_picked_small_context_deployment(proxy: ProxyClient, name: str)
)
@step("register zero-weight backup deployment")
def create_zero_weight_backup_deployment(proxy: ProxyClient, name: str) -> str:
"""The other half of a retry pair: healthy, but weight 0, so the weighted shuffle
never opens on it. It is reachable only once its sibling is out of the running,
@ -262,6 +277,7 @@ def create_zero_weight_backup_deployment(proxy: ProxyClient, name: str) -> str:
)
@step("POST /chat/completions")
def chat_turns_override(
proxy: ProxyClient,
key: str,
@ -289,6 +305,7 @@ def chat_turns_override(
)
@step("POST /chat/completions")
def chat_override(
proxy: ProxyClient,
key: str,
@ -311,6 +328,7 @@ def chat_override(
)
@step("POST /chat/completions (stream held open)")
def open_chat_stream(
proxy: ProxyClient,
key: str,

301
tests/e2e/test_junit_properties.py
View file

@ -10,12 +10,18 @@ rollups and, for ``source``, the status page's per-test links to GitHub.
from __future__ import annotations
import threading
import warnings
from collections.abc import Generator
from contextlib import contextmanager
from pathlib import Path
import pytest
from e2e_metadata import MAX_STEPS, STEP_FRAMES, STEPS, step, step_properties
from junit_properties import (
SUITE_ROOT,
attach_result_properties,
attach_step_properties,
dedupe_covers,
package_from_nodeid,
result_properties,
@ -129,3 +135,298 @@ class TestSuiteRoot:
class TestDedupeCovers:
def test_ids_are_unique_order_preserving_and_non_empty_strings(self) -> None:
assert dedupe_covers([("A", "B"), ("B", ""), ("C", 7)]) == ("A", "B", "C")
class TestStepRecording:
"""`@step`-decorated harness helpers append to the running test's story as
they execute.
Each test here starts from an empty log because conftest's
`pytest_runtest_setup` hook resets the recorder first thing in every test's
setup -- the same reset the live suite relies on for per-test isolation.
"""
def test_steps_land_in_call_order(self) -> None:
@step("register deployment")
def register() -> str:
return "model-id"
@step("generate virtual key")
def generate() -> str:
return "sk-x"
_ = register()
_ = generate()
assert STEPS.taken() == ("register deployment", "generate virtual key")
def test_a_decorated_helper_still_returns_exactly_what_it_did(self) -> None:
"""`@step` records, it does not intercept: arguments, return value and
`__name__` all survive it, so decorating a live harness method cannot
change what the test observes."""
@step("POST /chat/completions")
def chat(key: str, *, model: str) -> str:
return f"{key}:{model}"
assert chat("sk-x", model="gpt-5.5") == "sk-x:gpt-5.5"
assert chat.__name__ == "chat"
def test_a_helper_that_raises_leaves_its_own_label_last(self) -> None:
"""The whole point of the field. The label is recorded BEFORE the call, so
a test that dies inside a helper keeps a partial story whose last element
names the helper it died in."""
@step("generate virtual key")
def generate() -> str:
return "sk-x"
@step("POST /chat/completions")
def chat() -> None:
raise RuntimeError("502 from upstream")
_ = generate()
with pytest.raises(RuntimeError, match="502 from upstream"):
chat()
assert STEPS.taken() == ("generate virtual key", "POST /chat/completions")
def test_a_poll_loop_is_one_step_in_the_story_not_fifty(self) -> None:
@step("poll /spend/logs for the request id")
def poll() -> None:
return None
for _ in range(20):
poll()
assert STEPS.taken() == ("poll /spend/logs for the request id",)
def test_the_same_label_recorded_again_later_is_a_new_step(self) -> None:
"""Only CONSECUTIVE duplicates collapse; a helper called again after
something else happened is a genuine second beat of the story."""
STEPS.record("POST /chat/completions")
STEPS.record("poll /spend/logs")
STEPS.record("POST /chat/completions")
assert STEPS.taken() == ("POST /chat/completions", "poll /spend/logs", "POST /chat/completions")
def test_the_log_is_capped_so_a_load_test_cannot_bury_the_story(self) -> None:
for index in range(MAX_STEPS * 2):
STEPS.record(f"call {index}")
taken = STEPS.taken()
assert len(taken) == MAX_STEPS
assert taken[0] == "call 0"
def test_whitespace_is_normalized_and_an_empty_label_records_nothing(self) -> None:
STEPS.record(" POST /chat/completions\n ")
STEPS.record(" ")
assert STEPS.taken() == ("POST /chat/completions",)
def test_reset_empties_the_log_so_one_test_never_inherits_another_s(self) -> None:
STEPS.record("register deployment")
STEPS.reset()
assert STEPS.taken() == ()
assert step_properties() == ()
def test_steps_serialize_as_repeated_properties_in_order(self) -> None:
"""Repeated rather than joined on a delimiter: the labels are free text, so
no separator can be reserved, and a repeated property has none to corrupt."""
STEPS.record('attach guardrail, comma & "quoted" <tag>')
STEPS.record("POST /chat/completions")
assert step_properties() == (
("step", 'attach guardrail, comma & "quoted" <tag>'),
("step", "POST /chat/completions"),
)
def test_a_decorated_helper_warns_at_its_caller_with_step_frames(self) -> None:
"""`stacklevel` counts frames, and the wrapper is one of them: a cleanup
helper that warns about its caller would otherwise report every warning at
e2e_metadata.py. Pins `STEP_FRAMES` to the frames the wrapper really adds."""
@step("delete team")
def delete_team() -> None:
warnings.warn("delete_team('t') failed", stacklevel=2 + STEP_FRAMES)
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
delete_team()
assert [Path(warning.filename).name for warning in caught] == [Path(__file__).name]
class TestNestedSteps:
"""Harness layers call each other, so a step's helper routinely calls other
decorated helpers. Only the outermost records."""
def test_a_step_called_inside_a_step_is_not_recorded(self) -> None:
"""`ResourceManager.key` wraps `ProxyClient.generate_key`: one action, one
beat of the story, at the level the test called in at."""
@step("POST /key/generate")
def generate_key() -> str:
return "sk-x"
@step("generate virtual key")
def key() -> str:
return generate_key()
assert key() == "sk-x"
assert STEPS.taken() == ("generate virtual key",)
def test_the_inner_step_records_again_once_the_outer_one_returns(self) -> None:
@step("POST /key/generate")
def generate_key() -> str:
return "sk-x"
@step("generate virtual key")
def key() -> str:
return generate_key()
_ = key()
_ = generate_key()
assert STEPS.taken() == ("generate virtual key", "POST /key/generate")
def test_an_inner_step_that_raises_leaves_the_outer_label_last_and_unwinds(self) -> None:
"""The helper the test called is where it died, and the nesting flag is
released on the way out, so the next top-level call still records."""
@step("POST /team/new")
def post_team() -> None:
raise RuntimeError("/team/new answered 500")
@step("create team with a budget")
def create_team() -> None:
post_team()
@step("POST /chat/completions")
def chat() -> None:
return None
with pytest.raises(RuntimeError, match="answered 500"):
create_team()
chat()
assert STEPS.taken() == ("create team with a budget", "POST /chat/completions")
def test_a_worker_thread_a_step_fans_out_to_records_its_own_steps(self) -> None:
"""Nesting is per thread: a load helper that fans chats out to workers is
not inside a step on those workers, so their calls are still recorded."""
@step("POST /chat/completions")
def chat() -> None:
return None
@step("fire concurrent chats")
def fan_out() -> None:
worker = threading.Thread(target=chat)
worker.start()
worker.join()
fan_out()
assert STEPS.taken() == ("fire concurrent chats", "POST /chat/completions")
class TestContextManagerSteps:
"""A `@contextmanager` helper's setup and cleanup run at `__enter__` and
`__exit__`, after the decorated call has returned. Both still count as part
of its step; the `with` body is the test's own code and records as usual."""
def test_setup_and_cleanup_stay_inside_the_step_and_the_body_records(self) -> None:
@step("run a SQL statement")
def execute() -> None:
return None
@step("create a read-only database role")
@contextmanager
def restricted_user() -> Generator[str]:
execute()
try:
yield "reader"
finally:
execute()
@step("POST /chat/completions")
def chat() -> None:
return None
with restricted_user() as user:
assert user == "reader"
chat()
assert STEPS.taken() == ("create a read-only database role", "POST /chat/completions")
def test_a_test_that_dies_in_the_with_body_keeps_its_last_step_last(self) -> None:
"""The guarantee the field makes: the cleanup that runs on the way out of
the `with` must not append a step behind the one the test died on."""
@step("drop the role")
def drop_role() -> None:
return None
@step("create a read-only database role")
@contextmanager
def restricted_user() -> Generator[None]:
try:
yield
finally:
drop_role()
@step("POST /chat/completions")
def chat() -> None:
raise RuntimeError("502 from upstream")
with pytest.raises(RuntimeError, match="502 from upstream"), restricted_user():
chat()
assert STEPS.taken() == ("create a read-only database role", "POST /chat/completions")
def test_the_wrapped_context_keeps_its_exception_handling(self) -> None:
"""`__exit__` is forwarded, return value included, so a context that
suppresses an exception still does."""
@step("hold an advisory lock")
@contextmanager
def swallowing() -> Generator[None]:
try:
yield
except KeyError:
pass
with swallowing():
raise KeyError("suppressed by the context")
assert STEPS.taken() == ("hold an advisory lock",)
def test_a_bare_generator_is_refused_where_the_decorator_runs(self) -> None:
"""Its body runs only as the caller iterates, interleaved with the caller's
own steps, so no single point in the story is where it happened. Refused at
decoration, which for a harness module is import, so it lands as a
collection error rather than a story that quietly reads out of order."""
def rows() -> Generator[int]:
yield 1
with pytest.raises(TypeError, match="cannot wrap the generator function"):
_ = step("poll /spend/logs")(rows)
class TestAttachStepProperties:
def test_steps_are_appended_after_the_collected_properties(self, request: pytest.FixtureRequest) -> None:
"""Order inside `<properties>` is list order, so the story reads after the
fixed prefix the collection hook already attached."""
test = type(self).test_steps_are_appended_after_the_collected_properties
item = collected_item(request, test.__name__)
STEPS.record("register deployment")
STEPS.record("POST /chat/completions")
attach_step_properties(item)
assert [name for name, _ in item.user_properties] == ["package", "covers", "source", "step", "step"]
assert [value for name, value in item.user_properties if name == "step"] == [
"register deployment",
"POST /chat/completions",
]
def test_a_rerun_replaces_the_story_rather_than_appending_a_second_one(
self, request: pytest.FixtureRequest
) -> None:
"""The suite runs with `--reruns 1`. Without this the retry's steps would
queue up behind the first attempt's and the report would read as one test
that did everything twice."""
test = type(self).test_a_rerun_replaces_the_story_rather_than_appending_a_second_one
item = collected_item(request, test.__name__)
STEPS.record("attempt one died here")
attach_step_properties(item)
STEPS.reset()
STEPS.record("attempt two got further")
attach_step_properties(item)
assert [value for name, value in item.user_properties if name == "step"] == ["attempt two got further"]

288
tests/e2e/test_junit_report.py Normal file
View file

@ -0,0 +1,288 @@
"""The JUnit report itself, written by a real pytest run.
No proxy and no ``e2e`` marker. test_junit_properties.py pins the functions that
build the properties; this pins what reaches the XML once pytest, its junitxml
plugin, pytest-rerunfailures and xdist are all in the loop. Each case writes a
throwaway suite into a tmp dir and runs it in a child interpreter with THIS
directory's conftest.py loaded as a plugin, so the hooks under test are the ones
the live suite runs and the recorder is the real one, never a copy of either.
The timing that makes the recorded half work is pytest's, which is why it is
pinned here against the real thing: junitxml writes a testcase's properties from
its TEARDOWN report, and pytest builds that report from ``item.user_properties``
after the setup and call phases have both attached the steps. The suite runs
distributed, so every assertion is made in-process and again under ``-n 2``.
"""
from __future__ import annotations
import os
import shlex
import subprocess
import sys
from collections.abc import Mapping
from importlib.util import find_spec
from pathlib import Path
from types import MappingProxyType
from typing import Final
from xml.etree import ElementTree
import pytest
SUITE_DIR: Final = Path(__file__).resolve().parent
CHILD_TIMEOUT_SECONDS: Final = 180
STORY_SUITE: Final = """
from collections.abc import Iterator
from pathlib import Path
import pytest
from e2e_metadata import step
FIRST_ATTEMPT_MADE = Path(__file__).with_name("first-attempt-made")
@step("generate virtual key")
def generate_key() -> None:
return None
@step("create team")
def create_team() -> None:
raise RuntimeError("/team/new answered 500")
@step("POST /chat/completions")
def chat(*, ok: bool) -> None:
if not ok:
raise AssertionError("status_code=502 from upstream")
@step("poll /spend/logs")
def poll_spend_logs() -> None:
return None
@step("delete virtual key")
def delete_key() -> None:
return None
@pytest.fixture
def key() -> Iterator[None]:
generate_key()
yield
delete_key()
@pytest.fixture
def team(key: None) -> None:
create_team()
def test_passes(key: None) -> None:
chat(ok=True)
poll_spend_logs()
def test_fails(key: None) -> None:
chat(ok=False)
poll_spend_logs()
def test_errors_in_setup(team: None) -> None:
poll_spend_logs()
def test_passes_on_the_rerun(key: None) -> None:
first_attempt = not FIRST_ATTEMPT_MADE.exists()
FIRST_ATTEMPT_MADE.touch()
chat(ok=not first_attempt)
poll_spend_logs()
"""
WIDE_FINALIZER_SUITE: Final = """
from collections.abc import Iterator
import pytest
from e2e_metadata import step
@step("generate virtual key")
def generate_key() -> None:
return None
@step("delete shared team")
def delete_shared_team() -> None:
return None
@pytest.fixture(scope="module")
def shared_team() -> Iterator[None]:
yield
delete_shared_team()
def test_uses_the_shared_team(shared_team: None) -> None:
generate_key()
"""
WIDE_SETUP_ERROR_SUITE: Final = """
import pytest
from e2e_metadata import step
@step("log in to the identity provider")
def log_in() -> None:
raise RuntimeError("identity provider is down")
@pytest.fixture(scope="module")
def identity() -> None:
log_in()
def test_dies_in_a_module_scoped_fixture(identity: None) -> None:
assert identity is None
"""
Properties = tuple[tuple[str, str], ...]
def write_suite(directory: Path, modules: Mapping[str, str]) -> None:
"""Lay a child suite out in ``directory``, with an ini file of its own.
The ini pins the child's rootdir to the tmp dir wherever that lives, and its
``pythonpath`` is what makes this directory's conftest.py, and the harness
modules the child suite imports, importable under ``-I``.
"""
_ = (directory / "pytest.ini").write_text(f"[pytest]\npythonpath = {shlex.quote(str(SUITE_DIR))}\n")
for name, source in modules.items():
_ = (directory / name).write_text(source)
def run_child_pytest(suite: Path, *args: str) -> subprocess.CompletedProcess[str]:
"""Run pytest over ``suite`` in a fresh interpreter, hooked up like the live suite.
``-p conftest`` registers this directory's conftest.py as a plugin, since a
tmp dir outside tests/e2e would never pick it up by location. The parent's
fixture-mode and addopts settings are dropped so a replay lane cannot leak
into the child.
"""
inherited: Final = {
name: value
for name, value in os.environ.items()
if name != "PYTEST_ADDOPTS" and not name.startswith("E2E_FIXTURE_")
}
return subprocess.run(
[sys.executable, "-I", "-m", "pytest", "-p", "conftest", "-p", "no:cacheprovider", *args, str(suite)],
cwd=suite,
env=inherited,
capture_output=True,
text=True,
timeout=CHILD_TIMEOUT_SECONDS,
check=False,
)
def properties_by_test(testsuite: ElementTree.Element) -> Mapping[str, Properties]:
"""Every testcase's <property> pairs, in document order, keyed by test name."""
return MappingProxyType(
{
testcase.get("name", ""): tuple(
(prop.get("name", ""), prop.get("value", "")) for prop in testcase.iter("property")
)
for testcase in testsuite.iter("testcase")
}
)
def values(properties: Properties, name: str) -> tuple[str, ...]:
return tuple(value for prop, value in properties if prop == name)
@pytest.fixture(
scope="module",
params=[
pytest.param((), id="in-process"),
pytest.param(
("-n", "2"),
id="xdist",
marks=pytest.mark.skipif(find_spec("xdist") is None, reason="pytest-xdist is not installed"),
),
],
)
def report(request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory) -> Mapping[str, Properties]:
"""One child run per distribution mode, shared by every assertion below.
``--reruns 1`` and the ``--only-rerun`` pattern are the live suite's own
addopts. The two wide-scope modules sort ahead of the story, and next to each
other, so in-process the second one's setup runs right after the first one's
module-scoped finalizer.
"""
distribution: Final[tuple[str, ...]] = request.param # pyright: ignore[reportAny] # pytest types request.param as Any
suite: Final = tmp_path_factory.mktemp("suite")
write_suite(
suite,
{
"test_scope_a_finalizer.py": WIDE_FINALIZER_SUITE,
"test_scope_b_setup_error.py": WIDE_SETUP_ERROR_SUITE,
"test_story.py": STORY_SUITE,
},
)
xml: Final = suite / "report.xml"
child: Final = run_child_pytest(
suite, f"--junitxml={xml}", "--reruns", "1", "--only-rerun", "status_code=5[0-9][0-9]", *distribution
)
assert xml.exists(), f"the child run wrote no JUnit report:\n{child.stdout}\n{child.stderr}"
testsuite: Final = next(ElementTree.parse(xml).getroot().iter("testsuite"))
outcomes: Final = {name: testsuite.get(name) for name in ("tests", "failures", "errors", "skipped")}
assert outcomes == {"tests": "6", "failures": "1", "errors": "2", "skipped": "0"}, child.stdout
return properties_by_test(testsuite)
class TestStepsReachTheReport:
def test_a_passing_test_tells_its_story_in_call_order(self, report: Mapping[str, Properties]) -> None:
"""Fixture setup first, then the body. The finalizer's "delete virtual key"
is cleanup and is deliberately not part of the story."""
assert values(report["test_passes"], "step") == (
"generate virtual key",
"POST /chat/completions",
"poll /spend/logs",
)
def test_a_failing_test_s_last_step_is_where_it_died(self, report: Mapping[str, Properties]) -> None:
"""The reason the field exists. Nothing the test never reached is listed,
and no teardown step is appended behind the one it died on."""
assert values(report["test_fails"], "step") == ("generate virtual key", "POST /chat/completions")
def test_a_setup_error_keeps_the_steps_recorded_before_the_crash(self, report: Mapping[str, Properties]) -> None:
"""A fixture that raises never reaches the call phase, and setup is where
an e2e test most often dies (proxy not ready, key creation failing), so
the steps have to be attached after setup too."""
assert values(report["test_errors_in_setup"], "step") == ("generate virtual key", "create team")
def test_a_rerun_reports_only_the_attempt_junit_records(self, report: Mapping[str, Properties]) -> None:
"""The first attempt died on the chat call and the rerun got through. Steps
are attached twice per attempt, and none of that may show up as a doubled
or a stale story."""
assert values(report["test_passes_on_the_rerun"], "step") == (
"generate virtual key",
"POST /chat/completions",
"poll /spend/logs",
)
def test_a_setup_error_does_not_inherit_a_wider_finalizer_s_steps(self, report: Mapping[str, Properties]) -> None:
"""A module-scoped finalizer runs after the last test of its module, and
a module-scoped fixture is set up before any function-scoped one. The log
is emptied ahead of both, so the next test's setup error reports its own
steps and not "delete shared team"."""
assert values(report["test_uses_the_shared_team"], "step") == ("generate virtual key",)
assert values(report["test_dies_in_a_module_scoped_fixture"], "step") == ("log in to the identity provider",)
def test_steps_ride_behind_the_fixed_prefix(self, report: Mapping[str, Properties]) -> None:
"""`package`/`covers`/`source` are what Loki, Grafana and the status page
already read, on every outcome including a setup error."""
for name in ("test_passes", "test_fails", "test_errors_in_setup"):
assert tuple(prop for prop, _ in report[name])[:4] == ("package", "covers", "source", "step"), name