diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index b00b7dfac95..3bdc8af2f52 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -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 `` 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 diff --git a/tests/e2e/a2a/a2a_client.py b/tests/e2e/a2a/a2a_client.py index 605dd8fb7e5..7b2cdd339c4 100644 --- a/tests/e2e/a2a/a2a_client.py +++ b/tests/e2e/a2a/a2a_client.py @@ -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 diff --git a/tests/e2e/access_control/access_control_client.py b/tests/e2e/access_control/access_control_client.py index 5f459c09767..c9b14415f4f 100644 --- a/tests/e2e/access_control/access_control_client.py +++ b/tests/e2e/access_control/access_control_client.py @@ -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", diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py index 86e47c0b1e1..d7547b7ec95 100644 --- a/tests/e2e/batches/batch_cleanup.py +++ b/tests/e2e/batches/batch_cleanup.py @@ -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, diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 8745140a818..0c19e9d0421 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -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}", diff --git a/tests/e2e/claude_code/_basic_messaging.py b/tests/e2e/claude_code/_basic_messaging.py index 7c581cc5e38..90485ac8549 100644 --- a/tests/e2e/claude_code/_basic_messaging.py +++ b/tests/e2e/claude_code/_basic_messaging.py @@ -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, diff --git a/tests/e2e/claude_code/_passthrough.py b/tests/e2e/claude_code/_passthrough.py index be7a475dff7..5467eef087e 100644 --- a/tests/e2e/claude_code/_passthrough.py +++ b/tests/e2e/claude_code/_passthrough.py @@ -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, diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index a01d8ab3e7c..ff64219ae2c 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -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], diff --git a/tests/e2e/claude_code/http_probe.py b/tests/e2e/claude_code/http_probe.py index 8aba54576c4..5ef2981aab0 100644 --- a/tests/e2e/claude_code/http_probe.py +++ b/tests/e2e/claude_code/http_probe.py @@ -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, diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 8776d00d502..2a1dbb2ec01 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -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 diff --git a/tests/e2e/e2e_metadata.py b/tests/e2e/e2e_metadata.py new file mode 100644 index 00000000000..46911745871 --- /dev/null +++ b/tests/e2e/e2e_metadata.py @@ -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 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()) diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index 3ceac737399..03d6d50a43f 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -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 diff --git a/tests/e2e/idp.py b/tests/e2e/idp.py index a89a036baeb..96cbce970b4 100644 --- a/tests/e2e/idp.py +++ b/tests/e2e/idp.py @@ -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: diff --git a/tests/e2e/junit_properties.py b/tests/e2e/junit_properties.py index b9f5da871ae..9ee1ceebc96 100644 --- a/tests/e2e/junit_properties.py +++ b/tests/e2e/junit_properties.py @@ -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()) diff --git a/tests/e2e/lifecycle.py b/tests/e2e/lifecycle.py index eb9704d4dcb..b08579d3038 100644 --- a/tests/e2e/lifecycle.py +++ b/tests/e2e/lifecycle.py @@ -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])) diff --git a/tests/e2e/llm_translation/passthrough_client.py b/tests/e2e/llm_translation/passthrough_client.py index a56d3dc077e..7f2419fa67c 100644 --- a/tests/e2e/llm_translation/passthrough_client.py +++ b/tests/e2e/llm_translation/passthrough_client.py @@ -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, diff --git a/tests/e2e/llm_translation/realtime/realtime_client.py b/tests/e2e/llm_translation/realtime/realtime_client.py index 7c4a9cc4af9..382d347cddc 100644 --- a/tests/e2e/llm_translation/realtime/realtime_client.py +++ b/tests/e2e/llm_translation/realtime/realtime_client.py @@ -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 diff --git a/tests/e2e/load/session_anomaly.py b/tests/e2e/load/session_anomaly.py index c29b833635d..40b2d4fc07c 100644 --- a/tests/e2e/load/session_anomaly.py +++ b/tests/e2e/load/session_anomaly.py @@ -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, diff --git a/tests/e2e/logging/datadog_reader.py b/tests/e2e/logging/datadog_reader.py index 368c20cb6aa..aab7c0bc9b7 100644 --- a/tests/e2e/logging/datadog_reader.py +++ b/tests/e2e/logging/datadog_reader.py @@ -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), diff --git a/tests/e2e/logging/gcs_reader.py b/tests/e2e/logging/gcs_reader.py index 60622c121ac..d422d17429b 100644 --- a/tests/e2e/logging/gcs_reader.py +++ b/tests/e2e/logging/gcs_reader.py @@ -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 diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index 66dfa233ec4..4b6e92ea2d8 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -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 diff --git a/tests/e2e/logging/s3_reader.py b/tests/e2e/logging/s3_reader.py index d605dec6096..27f9de1e15e 100644 --- a/tests/e2e/logging/s3_reader.py +++ b/tests/e2e/logging/s3_reader.py @@ -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 diff --git a/tests/e2e/logging/weave_reader.py b/tests/e2e/logging/weave_reader.py index 2f8f759d299..c1f3312ba4a 100644 --- a/tests/e2e/logging/weave_reader.py +++ b/tests/e2e/logging/weave_reader.py @@ -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 diff --git a/tests/e2e/management/jwt_actors.py b/tests/e2e/management/jwt_actors.py index 2d23549fe71..a5f62c723b0 100644 --- a/tests/e2e/management/jwt_actors.py +++ b/tests/e2e/management/jwt_actors.py @@ -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: diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 8470d318db8..fa1040b672f 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -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) diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index 352b4446cfd..eeb1fb06eca 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -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, diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 56f7fffba29..45c2d21afbc 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -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, diff --git a/tests/e2e/mcp/oauth_chat_client.py b/tests/e2e/mcp/oauth_chat_client.py index 0c5c6106259..1e8268e6e3a 100644 --- a/tests/e2e/mcp/oauth_chat_client.py +++ b/tests/e2e/mcp/oauth_chat_client.py @@ -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 diff --git a/tests/e2e/mcp/oauth_gateway.py b/tests/e2e/mcp/oauth_gateway.py index 82bb5f7ba0b..1f2dc7e977c 100644 --- a/tests/e2e/mcp/oauth_gateway.py +++ b/tests/e2e/mcp/oauth_gateway.py @@ -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" diff --git a/tests/e2e/migrations/checks.py b/tests/e2e/migrations/checks.py index 619ad3b0e6c..30f206468d3 100644 --- a/tests/e2e/migrations/checks.py +++ b/tests/e2e/migrations/checks.py @@ -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( diff --git a/tests/e2e/migrations/containers.py b/tests/e2e/migrations/containers.py index dd126b994d3..41e10f90e70 100644 --- a/tests/e2e/migrations/containers.py +++ b/tests/e2e/migrations/containers.py @@ -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, diff --git a/tests/e2e/migrations/database.py b/tests/e2e/migrations/database.py index a370c21ba0b..272f413180b 100644 --- a/tests/e2e/migrations/database.py +++ b/tests/e2e/migrations/database.py @@ -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]}" diff --git a/tests/e2e/migrations/upgrade.py b/tests/e2e/migrations/upgrade.py index 2123f86450a..6b9ee0ba854 100644 --- a/tests/e2e/migrations/upgrade.py +++ b/tests/e2e/migrations/upgrade.py @@ -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" diff --git a/tests/e2e/otel_client.py b/tests/e2e/otel_client.py index b11fddebc9c..d07f59fabf3 100644 --- a/tests/e2e/otel_client.py +++ b/tests/e2e/otel_client.py @@ -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]: diff --git a/tests/e2e/other/other_client.py b/tests/e2e/other/other_client.py index 4313bbe4068..bc65b6a22a0 100644 --- a/tests/e2e/other/other_client.py +++ b/tests/e2e/other/other_client.py @@ -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 diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 76c57452c69..8e391ab69f4 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -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()) diff --git a/tests/e2e/quota_management/budgets/budget_client.py b/tests/e2e/quota_management/budgets/budget_client.py index 087dc8ca522..701d67f8c6e 100644 --- a/tests/e2e/quota_management/budgets/budget_client.py +++ b/tests/e2e/quota_management/budgets/budget_client.py @@ -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", diff --git a/tests/e2e/quota_management/ratelimit/quota_client.py b/tests/e2e/quota_management/ratelimit/quota_client.py index a3a467a1d71..4d0ccfc52e0 100644 --- a/tests/e2e/quota_management/ratelimit/quota_client.py +++ b/tests/e2e/quota_management/ratelimit/quota_client.py @@ -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", diff --git a/tests/e2e/quota_management/spend_tracking/cost_rows.py b/tests/e2e/quota_management/spend_tracking/cost_rows.py index 87af54fe83f..3bf3f6a90e7 100644 --- a/tests/e2e/quota_management/spend_tracking/cost_rows.py +++ b/tests/e2e/quota_management/spend_tracking/cost_rows.py @@ -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, diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index b7f59fe5f89..b8406217a9e 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -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( diff --git a/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py b/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py index 26809874aed..0b495cde29e 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py +++ b/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py @@ -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" diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index 3d5b76f6408..4fcf3ed40c2 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -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, diff --git a/tests/e2e/test_junit_properties.py b/tests/e2e/test_junit_properties.py index 02c1413c840..9b660470309 100644 --- a/tests/e2e/test_junit_properties.py +++ b/tests/e2e/test_junit_properties.py @@ -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" ') + STEPS.record("POST /chat/completions") + assert step_properties() == ( + ("step", 'attach guardrail, comma & "quoted" '), + ("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 `` 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"] diff --git a/tests/e2e/test_junit_report.py b/tests/e2e/test_junit_report.py new file mode 100644 index 00000000000..508c4e9f417 --- /dev/null +++ b/tests/e2e/test_junit_report.py @@ -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 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