scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

test(e2e): tag a2a, access_control, other, secret_manager and migrations tests with Subject metadata (#44965)

Browse source 
* test(e2e): add enum values, auto-discovering label gates and secret hiding for e2e metadata

* test(e2e): tag a2a, access_control, other, secret_manager and migrations tests with Subject metadata

* docs(e2e): name every markerless harness test file that carries no Subject

* test(e2e): keep the step discovery comprehensions to one for clause
This commit is contained in:
ryan-crabbe-berri 2026-10-07 12:11:55 -07:00 • committed by GitHub
parent aa3cd70c18
commit 40a9b959a6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
27 changed files with 578 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -19,6 +19,7 @@ from pydantic import BaseModel, ConfigDict, Field
from e2e_config import settle_propagation
from e2e_http import NoBody, Result, Success, get_external, is_ok
from e2e_metadata import STEP_FRAMES, step
from proxy_client import ProxyClient
@ -87,7 +88,7 @@ class A2ABridgeParams(BaseModel):
custom_llm_provider: str
model: str
api_key: str | None = None
api_key: str | None = Field(default=None, repr=False)
class AgentRegisterBody(BaseModel):
@ -291,6 +292,7 @@ class A2AResponse(BaseModel):
class A2AClient:
proxy: ProxyClient
@step("Register the A2A agent {body.agent_name} through /v1/agents")
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("Read the A2A agent back from /v1/agents/{{agent_id}}")
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 the 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("Read the A2A agent's card from /a2a/{{agent_id}}/.well-known/agent-card.json with the given key")
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("Send an A2A message to /a2a/{{agent_id}} with {body.params.message.parts}")
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 a published A2A agent card from its /.well-known endpoint")
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

73
tests/e2e/a2a/test_a2a_agent_e2e.py
View file

@ -10,6 +10,8 @@ protocol version, and an unsupported version is refused at registration).
from __future__ import annotations
from typing import Final
import pytest
from a2a_client import (
@ -31,6 +33,9 @@ from a2a_client import (
from e2e_config import unique_marker
from e2e_http import Result, UnknownApiError, unwrap
from lifecycle import ResourceManager
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
BRIDGE_MODEL: Final = "claude-haiku-4-5"
# No api_key: litellm resolves ANTHROPIC_API_KEY from the proxy's own environment
# for this provider, which is what the agent-owner flow relies on. Pinning
@ -42,7 +47,7 @@ from lifecycle import ResourceManager
# omitted -> 200, "os.environ/..." -> 500 invalid x-api-key, literal key -> 200.
BRIDGE = A2ABridgeParams(
custom_llm_provider="anthropic",
model="claude-haiku-4-5",
model=BRIDGE_MODEL,
)
MOVEHOME_AGENT_CARD_URL = "https://movehome.org/.well-known/agent.json"
@ -96,6 +101,12 @@ def _ask(text: str) -> A2AJsonRpcRequest:
class TestA2AAgentLifecycle:
@pytest.mark.covers("other.a2a.register.persists")
@meta(
Subject(
domain=Domain.AGENTS_API,
route=Route.A2A,
)
)
def test_register_persists(self, client: A2AClient, resources: ResourceManager) -> None:
agent = _register(client, resources, "0.3")
fetched = unwrap(client.get_agent(agent.agent_id))
@ -104,6 +115,15 @@ class TestA2AAgentLifecycle:
assert fetched.agent_card_params.protocol_version == "0.3"
@pytest.mark.covers("other.a2a.register.semver_version_accepted")
@meta(
Subject(
domain=Domain.AGENTS_API,
route=Route.A2A,
providers=(Provider.ANTHROPIC,),
models=(BRIDGE_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_semver_protocol_version_registers_and_serves(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None:
agent = _register(client, resources, "0.3.0")
assert agent.agent_card_params.protocol_version == "0.3"
@ -116,6 +136,12 @@ class TestA2AAgentLifecycle:
assert result.text != ""
@pytest.mark.covers("other.a2a.message_send.real_world_agent_replies")
@meta(
Subject(
domain=Domain.AGENTS_API,
route=Route.A2A,
)
)
def test_real_world_agent_replies_to_property_query(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None:
upstream = unwrap(fetch_agent_card(MOVEHOME_AGENT_CARD_URL)).model_copy(update={"url": MOVEHOME_ORIGIN})
assert upstream.protocol_version == "0.3.0"
@ -152,6 +178,12 @@ class TestA2AAgentLifecycle:
assert all(listing.location.un_locode == location for listing in results.listings)
@pytest.mark.covers("other.a2a.discovery.proxy_fronted_card")
@meta(
Subject(
domain=Domain.AGENTS_API,
route=Route.A2A,
)
)
def test_discovery_card_is_proxy_fronted(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None:
agent = _register(client, resources, "0.3")
card = unwrap(client.agent_card(agent.agent_id, scoped_key))
@ -163,6 +195,15 @@ class TestA2AAgentLifecycle:
assert card.supported_interfaces[0].url == card.url
@pytest.mark.covers("other.a2a.message_send.bridge_invokes")
@meta(
Subject(
domain=Domain.AGENTS_API,
route=Route.A2A,
providers=(Provider.ANTHROPIC,),
models=(BRIDGE_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_message_send_runs_completion_bridge(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None:
agent = _register(client, resources, "0.3")
request = _ask("Reply with exactly the word PONG and nothing else")
@ -177,6 +218,15 @@ class TestA2AAgentLifecycle:
assert rows[0].model == f"a2a_agent/{agent.agent_card_params.name}"
@pytest.mark.covers("other.a2a.version.serves_pinned_0_3")
@meta(
Subject(
domain=Domain.AGENTS_API,
route=Route.A2A,
providers=(Provider.ANTHROPIC,),
models=(BRIDGE_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_pinned_v0_3_serves_flat_message_shape(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None:
agent = _register(client, resources, "0.3")
request = _ask("Say hi in one word")
@ -188,6 +238,15 @@ class TestA2AAgentLifecycle:
assert result.text != ""
@pytest.mark.covers("other.a2a.version.serves_pinned_1_0")
@meta(
Subject(
domain=Domain.AGENTS_API,
route=Route.A2A,
providers=(Provider.ANTHROPIC,),
models=(BRIDGE_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_pinned_v1_0_serves_nested_message_shape(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None:
agent = _register(client, resources, "1.0")
request = _ask("Say hi in one word")
@ -199,6 +258,12 @@ class TestA2AAgentLifecycle:
assert result.text != ""
@pytest.mark.covers("other.a2a.register.unsupported_version_rejected")
@meta(
Subject(
domain=Domain.AGENTS_API,
route=Route.A2A,
)
)
def test_unsupported_protocol_version_rejected(self, client: A2AClient) -> None:
result = _register_rejection(client, "9.9")
match result:
@ -209,6 +274,12 @@ class TestA2AAgentLifecycle:
pytest.fail(f"expected 400 for unsupported protocolVersion, got {result}")
@pytest.mark.covers("other.a2a.register.malformed_version_rejected")
@meta(
Subject(
domain=Domain.AGENTS_API,
route=Route.A2A,
)
)
def test_malformed_protocol_version_rejected(self, client: A2AClient) -> None:
result = _register_rejection(client, "0.3.garbage")
match result:

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

@ -8,6 +8,7 @@ from dataclasses import dataclass
from pydantic import BaseModel, ValidationError
from proxy_client import ProxyClient
from e2e_metadata import step
from e2e_http import NoBody, StreamingResponse, is_ok, unwrap
from models import (
ChatBody,
@ -59,14 +60,17 @@ def error_envelope(body: str) -> ApiErrorEnvelope | None:
class AccessControlClient:
proxy: ProxyClient
@step("Generate a virtual key that can only call LLM API routes")
def llm_only_key(self) -> str:
return self.proxy.generate_key(
KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"])
)
@step("Delete the virtual key")
def delete_key(self, key: str) -> None:
self.proxy.delete_key(key)
@step('Send a /chat/completions request to {model} with the prompt "{content}"')
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 the team {team_alias} with models: {models}")
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 {team_alias}'s models to {models} through /team/update")
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 the 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("List the deployments in the model access group {access_group}")
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 from /team/info")
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("Add a deployment named {model_name} that calls openai/gpt-4o-mini with the given key")
def create_model_status(self, key: str, model_name: str) -> StreamingResponse:
return self.proxy.transport.send(
"/model/new",

43
tests/e2e/access_control/test_access_control_e2e.py
View file

@ -26,6 +26,7 @@ from e2e_http import Success, UnauthorizedError, UnknownApiError, unwrap
from lifecycle import ResourceManager
from models import ChatBody, ChatMessage, ChatResponse, EmbedBody, LiteLLMParamsBody
from proxy_client import ProxyClient
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
pytestmark = pytest.mark.e2e
@ -36,6 +37,14 @@ EMBEDDING_MODEL = "openai-text-embedding-3-small"
class TestAccessControl:
@meta(
Subject(
domain=Domain.PROXY_AUTH,
providers=(Provider.GEMINI,),
models=(ALLOWED_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_allowed_model_is_permitted(
self, client: AccessControlClient, resources: ResourceManager
) -> None:
@ -57,6 +66,13 @@ class TestAccessControl:
f"200 must carry a real completion, not an error envelope: {result.body[:300]}"
)
@meta(
Subject(
domain=Domain.PROXY_AUTH,
models=(DISALLOWED_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_disallowed_model_is_denied_403(
self, client: AccessControlClient, resources: ResourceManager
) -> None:
@ -73,6 +89,13 @@ class TestAccessControl:
)
@pytest.mark.covers("other.auth.virtual_key.route_group_allowed")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
providers=(Provider.GEMINI, Provider.OPENAI,),
models=(ALLOWED_MODEL, EMBEDDING_MODEL,),
)
)
def test_llm_api_routes_group_grants_every_llm_endpoint(
self, client: AccessControlClient, resources: ResourceManager
) -> None:
@ -97,6 +120,12 @@ class TestAccessControl:
f"the same key must still be shut out of /model/new, got {denied.status_code}: {denied.body[:300]}"
)
@meta(
Subject(
domain=Domain.PROXY_AUTH,
route=Route.MODEL_MANAGEMENT,
)
)
def test_llm_only_key_forbidden_from_management_route_403(
self, client: AccessControlClient, resources: ResourceManager
) -> None:
@ -111,6 +140,12 @@ class TestAccessControl:
f"403 body must be a route-permission denial, got: {result.body[:300]}"
)
@meta(
Subject(
domain=Domain.PROXY_AUTH,
mode=Mode.NONSTREAM,
)
)
def test_unknown_model_returns_400(
self, client: AccessControlClient, resources: ResourceManager
) -> None:
@ -140,6 +175,14 @@ class TestVirtualKeyAuth:
"mgmt.virtual_key.invalid_denied",
exercised_on=[],
)
@meta(
Subject(
domain=Domain.PROXY_AUTH,
providers=(Provider.ANTHROPIC,),
models=(VIRTUAL_KEY_BACKEND,),
mode=Mode.NONSTREAM,
)
)
def test_valid_key_allows_and_invalid_key_denied(
self, proxy: ProxyClient, resources: ResourceManager
) -> None:

31
tests/e2e/access_control/test_chat_auth_headers_e2e.py
View file

@ -11,6 +11,7 @@ import pytest
from e2e_http import AuthHeaders, NoBody, StreamingResponse, assert_auth_denied
from models import ChatBody, ChatMessage
from proxy_client import ProxyClient
from e2e_metadata import Domain, Route, Subject, meta
pytestmark = pytest.mark.e2e
@ -32,26 +33,56 @@ def _chat_with_headers(proxy: ProxyClient, headers: AuthHeaders | NoBody) -> Str
class TestChatAuthHeaders:
@pytest.mark.covers("other.auth.llm_chat.missing_header_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
route=Route.CHAT_COMPLETIONS,
)
)
def test_missing_authorization_header_is_denied(self, proxy: ProxyClient) -> None:
result = _chat_with_headers(proxy, NoBody())
assert_auth_denied(result, "missing Authorization")
@pytest.mark.covers("other.auth.llm_chat.invalid_bearer_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
route=Route.CHAT_COMPLETIONS,
)
)
def test_bearer_invalid_token_is_denied(self, proxy: ProxyClient) -> None:
result = _chat_with_headers(proxy, AuthHeaders(authorization="Bearer invalid_token"))
assert_auth_denied(result, "Bearer invalid_token")
@pytest.mark.covers("other.auth.llm_chat.no_bearer_prefix_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
route=Route.CHAT_COMPLETIONS,
)
)
def test_token_without_bearer_prefix_is_denied(self, proxy: ProxyClient) -> None:
result = _chat_with_headers(proxy, AuthHeaders(authorization="invalid_token"))
assert_auth_denied(result, "token without Bearer prefix")
@pytest.mark.covers("other.auth.llm_chat.empty_bearer_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
route=Route.CHAT_COMPLETIONS,
)
)
def test_empty_bearer_token_is_denied(self, proxy: ProxyClient) -> None:
result = _chat_with_headers(proxy, AuthHeaders(authorization="Bearer "))
assert_auth_denied(result, "empty Bearer token")
@pytest.mark.covers("other.auth.llm_chat.not_bearer_scheme_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
route=Route.CHAT_COMPLETIONS,
)
)
def test_not_bearer_scheme_is_denied(self, proxy: ProxyClient) -> None:
result = _chat_with_headers(proxy, AuthHeaders(authorization="NotBearer validtoken123"))
assert_auth_denied(result, "NotBearer scheme")

30
tests/e2e/access_control/test_model_access_group_e2e.py
View file

@ -33,6 +33,7 @@ from models import (
ModelNewBody,
TeamInfoResponse,
)
from e2e_metadata import Domain, Mode, Provider, Subject, meta
pytestmark = pytest.mark.e2e
@ -177,6 +178,14 @@ class TestKeyScopedToAccessGroup:
"other.auth.model_access_group.member_allowed",
)
@pytest.mark.parametrize(("case", "select_model"), ALLOWED)
@meta(
Subject(
domain=Domain.PROXY_AUTH,
providers=(Provider.OPENAI,),
models=(GROUP_BACKEND, WILDCARD_BARE_MODEL, WILDCARD_PREFIXED_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_group_grants_every_deployment_in_it(
self,
case: str,
@ -202,6 +211,13 @@ class TestKeyScopedToAccessGroup:
@pytest.mark.covers("other.auth.model_access_group.non_member_denied")
@pytest.mark.parametrize(("case", "select_model"), DENIED)
@meta(
Subject(
domain=Domain.PROXY_AUTH,
models=(GROUP_BACKEND, UNCOVERED_OPENAI_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_group_grants_nothing_outside_it(
self,
case: str,
@ -228,6 +244,14 @@ class TestKeyScopedToAccessGroup:
class TestTeamScopedToAccessGroup:
@pytest.mark.covers("other.auth.model_access_group.team_wildcard_bare_name_allowed")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
providers=(Provider.OPENAI,),
models=(TEAM_WILDCARD_BARE_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_group_grants_the_teams_own_wildcard(
self, client: AccessControlClient, team_grant: TeamGrant
) -> None:
@ -248,6 +272,12 @@ class TestTeamScopedToAccessGroup:
)
@pytest.mark.covers("other.auth.model_access_group.team_non_member_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
mode=Mode.NONSTREAM,
)
)
def test_group_grants_the_team_nothing_outside_it(
self, client: AccessControlClient, team_grant: TeamGrant
) -> None:

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

@ -3,6 +3,7 @@ from contextlib import ExitStack
from typing import Final
from uuid import uuid4
from e2e_metadata import step
from psycopg import sql
from .containers import Containers, Replica, failed, until
@ -21,12 +22,14 @@ GATED: Final = Migration(
)
@step("Start {count} proxy containers on the test database")
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("Check that the migration {migration.name} ran 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("Apply the test migration by hand and record it in _prisma_migrations")
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("Check that the original _prisma_migrations row and its effect survived, with the row marked finished: {finished}")
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("Install a trigger that pauses the migration before it is marked finished")
def pause_completion(database: Database) -> None:
database.execute(
sql.SQL(
@ -67,6 +73,10 @@ def pause_completion(database: Database) -> None:
)
@step(
"Start a proxy container on the migration and kill it at its crash point, "
"with the migration SQL committed: {after_commit}"
)
def interrupt_owner(
containers: Containers, database: Database, after_commit: bool, *, stop_database_session: bool = True
) -> None:
@ -104,6 +114,7 @@ def interrupt_owner(
)
@step("Wait for every proxy container to refuse an unconfirmed migration and log recovery guidance")
def unconfirmed(replicas: tuple[Replica, ...], database: Database) -> None:
failed(replicas, "Migration completion could not be verified")
started: Final = str(

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

@ -11,6 +11,7 @@ from typing import Final
from uuid import uuid4
from e2e_http import NoBody, Success, unwrap
from e2e_metadata import step
from models import KeyGenerateBody, KeyGenerateResponse, KeyInfoParams, KeyInfoResponse
from transport import HttpTransport
@ -20,12 +21,14 @@ from .startup_models import ContainerState, Migration, Observation, Readiness
MASTER_KEY: Final = "sk-migration-ci-fixture"
@step("Run a docker command")
def docker(*args: str) -> str:
result: Final = subprocess.run(("docker", *args), capture_output=True, text=True, timeout=90)
assert result.returncode == 0, f"Docker operation failed: {result.stderr}"
return result.stdout.strip()
@step("Wait for {description}")
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("Read the proxy container's state from docker inspect")
def state(self) -> ContainerState:
return ContainerState.model_validate_json(docker("inspect", "--format", "{{json .State}}", self.name))
@step("Check whether the proxy container is running and ready on /health/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 proxy container's 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 proxy container")
def kill(self) -> None:
if self.state().Running:
docker("kill", self.name)
@step("Generate a virtual key on the proxy container and read it back from /key/info and the database")
def usable(self, database: Database) -> None:
alias: Final = f"migration-{uuid4().hex}"
key: Final = unwrap(
@ -86,6 +94,7 @@ class Replica:
) == ((alias,),)
@step("Wait for every proxy container to be ready, then generate and read back a virtual key on each")
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("Wait for the seed proxy container to be ready and finish building its request-log indexes")
def seeded(seed: Replica, database: Database) -> None:
ready((seed,), database)
until("the seed replica to finish its request-log indexes", lambda: request_log_indexes_built(database))
@ -110,6 +120,7 @@ def request_log_indexes_built(database: Database) -> bool:
) == ((2,),)
@step('Wait for every proxy container to refuse to start, logging "{marker}"')
def failed(replicas: tuple[Replica, ...], marker: str) -> None:
def all_stopped() -> bool:
observations: Final = tuple(replica.observe() for replica in replicas)
@ -122,6 +133,7 @@ def failed(replicas: tuple[Replica, ...], marker: str) -> None:
assert marker in replica.logs(), f"Startup failed outside the expected migration: {marker}"
@step("Check that every proxy container keeps waiting without serving for {seconds}s")
def waiting(replicas: tuple[Replica, ...], seconds: float) -> None:
deadline: Final = time.monotonic() + seconds
while time.monotonic() < deadline:
@ -139,6 +151,7 @@ class Containers:
def using(self, image: str) -> "Containers":
return replace(self, image=image)
@step("Start a proxy container on the test database")
@contextmanager
def start(
self,
@ -213,6 +226,7 @@ class Containers:
subprocess.run(("docker", "rm", "-f", name), capture_output=True, text=True, timeout=30, check=True)
@step("Write the migration {migration.name} into the proxy container's migration directory")
def write_migration(directory: Path, migration: Migration) -> None:
path: Final = directory / "prisma" / "migrations" / migration.name
path.mkdir(parents=True)

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

@ -11,6 +11,8 @@ import psycopg
from psycopg import sql
from pydantic import TypeAdapter
from e2e_metadata import step
Scalar = str | int | bool | None
ROWS: Final = TypeAdapter(tuple[tuple[Scalar, ...], ...])
GATE_KEY: Final = 39178002
@ -35,6 +37,7 @@ class Database:
container_url: str
schema: str = "public"
@step("Open a connection to the test database")
@contextmanager
def connection(self) -> Generator[psycopg.Connection[tuple[object, ...]]]:
with psycopg.connect(self.url, autocommit=True, connect_timeout=5) as connection:
@ -42,19 +45,23 @@ class Database:
connection.execute("SET statement_timeout = '15s'")
yield connection
@step("Run a SQL statement on the test database")
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 test 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 {name} exists in the test database")
def exists(self, name: str) -> bool:
return self.query("SELECT to_regclass(%s) IS NOT NULL", (name,)) == ((True,),)
@step("Read the migration history from _prisma_migrations")
def history(self) -> tuple[tuple[Scalar, ...], ...]:
if not self.exists("_prisma_migrations"):
return ()
@ -63,6 +70,7 @@ class Database:
"applied_steps_count, logs FROM _prisma_migrations ORDER BY id"
)
@step("List the database sessions 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 +79,7 @@ class Database:
(key >> 32, key & 0xFFFFFFFF),
)
@step("Hold an advisory lock on the test database")
@contextmanager
def lock(self, key: int = GATE_KEY) -> Generator[None]:
with self.connection() as connection:
@ -86,6 +95,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 +115,7 @@ class Databases:
connection.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name)))
@step("Create a read-only database role on the test database")
@contextmanager
def restricted_user(database: Database) -> Generator[Database]:
role: Final = f"migration_reader_{uuid4().hex[:16]}"

16
tests/e2e/migrations/test_legacy.py
View file

@ -7,6 +7,7 @@ import pytest
from .checks import COMPLETE, assert_completed, confirmed_history, assert_original_proof, start_replicas
from .containers import Containers, failed, ready, seeded
from .database import Database, Databases
from e2e_metadata import Domain, Subject, meta
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
@ -42,10 +43,20 @@ def adopt_legacy(containers: Containers, database: Database) -> None:
class TestLegacyMigrations:
@meta(
Subject(
domain=Domain.DB,
)
)
def test_matching_schema_warns_and_starts(self, containers: Containers, database: Database) -> None:
adopt_legacy(containers, database)
@pytest.mark.parametrize("fault", ("schema_drift", "custom_migrations", "empty_ledger"))
@meta(
Subject(
domain=Domain.DB,
)
)
def test_unrecognized_legacy_state_is_not_baselined(
self, containers: Containers, database: Database, fault: str
) -> None:
@ -64,6 +75,11 @@ class TestLegacyMigrations:
) == ((0,),)
@pytest.mark.parametrize("scenario", ("upgrade", "recovery", "legacy"))
@meta(
Subject(
domain=Domain.DB,
)
)
def test_non_default_schema(
self, containers: Containers, databases: Databases, scenario: Literal["upgrade", "recovery", "legacy"]
) -> None:

6
tests/e2e/migrations/test_pooling.py
View file

@ -13,6 +13,7 @@ from psycopg import sql
from .checks import COMPLETE, assert_completed
from .containers import Containers, docker, ready, until
from .database import Database, Databases, prisma_url, restricted_user
from e2e_metadata import Domain, Subject, meta
POOL_IMAGE: Final = (
"ghcr.io/cloudnative-pg/pgbouncer@sha256:e6ddfe22d845e603825e235dd8334b21ecd125abea2a2172478f556b8dee2bb8"
@ -94,6 +95,11 @@ def pool(database: Database, output: Path) -> Generator[str]:
class TestMigrationPooling:
@pytest.mark.parametrize("scenario,replica_count", (("fresh", 3), ("upgrade", 3), ("legacy", 3), ("upgrade", 6)))
@meta(
Subject(
domain=Domain.DB,
)
)
def test_direct_migrations_with_one_application_backend(
self,
containers: Containers,

36
tests/e2e/migrations/test_recovery.py
View file

@ -20,12 +20,18 @@ from .checks import (
from .containers import Containers, failed, ready, until, waiting
from .database import COORDINATOR_LOCK, GATE_KEY, Database
from .startup_models import Migration
from e2e_metadata import Domain, Subject, meta
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
class TestMigrationRecovery:
@pytest.mark.parametrize("after_commit", (False, True))
@meta(
Subject(
domain=Domain.DB,
)
)
def test_container_owner_crash(self, containers: Containers, database: Database, after_commit: bool) -> None:
interrupt_owner(containers, database, after_commit, stop_database_session=False)
history: Final = database.history()
@ -42,6 +48,11 @@ class TestMigrationRecovery:
assert database.history() == history
@pytest.mark.parametrize("after_commit", (False, True))
@meta(
Subject(
domain=Domain.DB,
)
)
def test_owner_and_database_session_crash(
self, containers: Containers, database: Database, after_commit: bool
) -> None:
@ -60,6 +71,11 @@ class TestMigrationRecovery:
assert database.history() == history
@pytest.mark.parametrize("later_failure", (False, True))
@meta(
Subject(
domain=Domain.DB,
)
)
def test_remaining_migrations_after_recovery(
self, containers: Containers, database: Database, later_failure: bool
) -> None:
@ -98,6 +114,11 @@ class TestMigrationRecovery:
assert database.query("SELECT id FROM migration_next") == ((2,),)
assert_original_proof(database, original, True)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_second_crash_during_recovery_is_atomic(self, containers: Containers, database: Database) -> None:
original: Final = confirmed_history(database)
pause_completion(database)
@ -114,6 +135,11 @@ class TestMigrationRecovery:
ready(start_replicas(stack, containers, database, (COMPLETE,)), database)
assert_original_proof(database, original, True)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_competing_recovery_rechecks_stale_failures(self, containers: Containers, database: Database) -> None:
original: Final = confirmed_history(database)
with ExitStack() as stack:
@ -132,6 +158,11 @@ class TestMigrationRecovery:
@pytest.mark.parametrize(
"fault", ("no_steps", "extra_steps", "failure_logs", "checksum", "duplicate_history", "missing_script")
)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_unproven_history_is_never_repaired(
self,
containers: Containers,
@ -172,6 +203,11 @@ class TestMigrationRecovery:
assert database.history() == history
assert database.query("SELECT id FROM migration_effect") == ((1,),)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_coordinator_timeout_preserves_proof(self, containers: Containers, database: Database) -> None:
original: Final = confirmed_history(database)
with database.lock(COORDINATOR_LOCK):

11
tests/e2e/migrations/test_rolling_upgrade.py
View file

@ -14,11 +14,17 @@ from .upgrade import (
migration_names,
provision,
)
from e2e_metadata import Domain, Subject, meta
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
class TestRollingUpgrade:
@meta(
Subject(
domain=Domain.DB,
)
)
def test_baseline_replica_keeps_serving_while_the_candidate_migrates(
self, containers: Containers, baseline_image: str, baseline_database: Database
) -> None:
@ -38,6 +44,11 @@ class TestRollingUpgrade:
assert CACHED_PLAN not in old.logs(), "The baseline replica hit a stale prepared statement"
assert old.state().Running, "The baseline replica died during the upgrade"
@meta(
Subject(
domain=Domain.DB,
)
)
def test_both_releases_serve_and_share_keys_during_the_overlap(
self, containers: Containers, baseline_image: str, baseline_database: Database
) -> None:

6
tests/e2e/migrations/test_shaped_database.py
View file

@ -5,6 +5,7 @@ import pytest
from .containers import Containers, ready
from .database import Database
from .upgrade import assert_history_clean, assert_upgraded, confirm, migration_names, provision
from e2e_metadata import Domain, Subject, meta
SPEND_ROWS: Final = 20_000
@ -22,6 +23,11 @@ def seed_spend_logs(database: Database, rows: int) -> None:
class TestPopulatedDatabaseUpgrade:
@meta(
Subject(
domain=Domain.DB,
)
)
def test_upgrade_completes_and_preserves_a_populated_spend_log(
self, containers: Containers, baseline_image: str, baseline_database: Database
) -> None:

46
tests/e2e/migrations/test_startup.py
View file

@ -7,12 +7,18 @@ from .checks import COMPLETE, FATAL, GATED, assert_completed, start_replicas
from .containers import Containers, failed, ready, until, waiting
from .database import PRISMA_LOCK, Database, Databases, restricted_user
from .startup_models import Migration
from e2e_metadata import Domain, Subject, meta
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
class TestMigrationStartup:
@pytest.mark.parametrize("replicas,v2", ((1, True), (3, True), (1, False)))
@meta(
Subject(
domain=Domain.DB,
)
)
def test_fresh_database(self, containers: Containers, databases: Databases, replicas: int, v2: bool) -> None:
with databases.create() as database, ExitStack() as stack:
ready(tuple(stack.enter_context(containers.start(database, v2=v2)) for _ in range(replicas)), database)
@ -21,11 +27,21 @@ class TestMigrationStartup:
) == ((0,),)
assert database.query("SELECT count(*) > 0 FROM _prisma_migrations") == ((True,),)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_concurrent_upgrade(self, containers: Containers, database: Database) -> None:
with ExitStack() as stack:
ready(start_replicas(stack, containers, database, (COMPLETE,)), database)
assert_completed(database)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_waiters_survive_prolonged_contention(self, containers: Containers, database: Database) -> None:
with ExitStack() as stack:
with database.lock():
@ -37,6 +53,11 @@ class TestMigrationStartup:
ready((owner, *followers), database)
assert_completed(database, GATED)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_lock_deadline_then_restart(self, containers: Containers, database: Database) -> None:
history: Final = database.history()
with database.lock(PRISMA_LOCK):
@ -51,6 +72,11 @@ class TestMigrationStartup:
ready((restarted,), database)
assert_completed(database)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_fatal_sql(self, containers: Containers, database: Database) -> None:
with ExitStack() as stack:
replicas: Final = start_replicas(stack, containers, database, (FATAL,))
@ -61,6 +87,11 @@ class TestMigrationStartup:
(COMPLETE.name, "%MIGRATION_TEST_FATAL%"),
) == ((1,),)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_duplicate_object_does_not_hide_incomplete_sql(self, containers: Containers, database: Database) -> None:
database.execute(
"CREATE TABLE migration_existing (id int PRIMARY KEY); INSERT INTO migration_existing VALUES (42)"
@ -77,6 +108,11 @@ class TestMigrationStartup:
) == ((True,),)
@pytest.mark.parametrize("v2", (True, False))
@meta(
Subject(
domain=Domain.DB,
)
)
def test_restart_preserves_history_and_data(self, containers: Containers, database: Database, v2: bool) -> None:
history: Final = database.history()
before: Final = database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token')
@ -86,12 +122,22 @@ class TestMigrationStartup:
assert database.history() == history
assert set(before).issubset(database.query('SELECT token FROM "LiteLLM_VerificationToken" ORDER BY token'))
@meta(
Subject(
domain=Domain.DB,
)
)
def test_disabled_migrations(self, containers: Containers, database: Database) -> None:
history: Final = database.history()
with containers.start(database, (FATAL,), disabled=True) as replica:
ready((replica,), database)
assert database.history() == history
@meta(
Subject(
domain=Domain.DB,
)
)
def test_insufficient_privileges(self, containers: Containers, database: Database) -> None:
history: Final = database.history()
with restricted_user(database) as limited:

16
tests/e2e/migrations/test_upgrade.py
View file

@ -7,11 +7,17 @@ from .checks import start_replicas
from .containers import Containers, ready
from .database import Database
from .upgrade import assert_history_clean, assert_upgraded, confirm, migration_names, provision
from e2e_metadata import Domain, Subject, meta
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
class TestReleaseUpgrade:
@meta(
Subject(
domain=Domain.DB,
)
)
def test_candidate_applies_the_pending_release_migrations(
self, containers: Containers, baseline_database: Database
) -> None:
@ -21,6 +27,11 @@ class TestReleaseUpgrade:
assert_upgraded(before, migration_names(baseline_database))
assert_history_clean(baseline_database)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_upgrade_preserves_keys_minted_by_the_baseline_release(
self, containers: Containers, baseline_image: str, baseline_database: Database
) -> None:
@ -34,6 +45,11 @@ class TestReleaseUpgrade:
assert_upgraded(before, migration_names(baseline_database))
confirm(new, key, alias)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_concurrent_replicas_upgrade_a_baseline_database_once(
self, containers: Containers, baseline_database: Database
) -> None:

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

@ -8,6 +8,7 @@ from typing import Final
from uuid import uuid4
from e2e_http import Result, Success, unwrap
from e2e_metadata import step
from models import (
KeyGenerateBody,
KeyGenerateResponse,
@ -24,6 +25,7 @@ from .database import Database
CACHED_PLAN: Final = "cached plan must not change result type"
@step("Generate a virtual key on the proxy container")
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("Check that the key {alias} resolves on the proxy container through /key/info")
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("Send /v1/models requests with the virtual key to proxy container {replica.name} in the background")
@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("Wait for {calls} more successful /v1/models calls from {description}")
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 the applied migration names from _prisma_migrations")
def migration_names(database: Database) -> frozenset[str]:
return frozenset(str(row[0]) for row in database.query("SELECT migration_name FROM _prisma_migrations"))
@step("Check that _prisma_migrations holds no unfinished, rolled-back or duplicated migration")
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"

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

@ -17,6 +17,7 @@ from dataclasses import dataclass
from typing import Final
from e2e_http import AnthropicHeaders, AuthHeaders, NoBody, ProbeResult, Result
from e2e_metadata import step
from idp import Keycloak, keycloak_from_env
from models import (
ChatBody,
@ -56,11 +57,13 @@ class OtherClient:
"""Resolved per use, so the suite's non-JWT tests never need the IdP env."""
return keycloak_from_env()
@step("Call /health/liveliness without credentials")
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("Call /health/readiness without credentials")
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."""
@ -71,6 +74,7 @@ class OtherClient:
response_type=ReadinessResponse,
)
@step("Call /health/readiness/details with the given key")
def readiness_details(self, key: str) -> Result[ReadinessDetailsResponse]:
return self.proxy.transport.get(
"/health/readiness/details",
@ -79,6 +83,7 @@ class OtherClient:
response_type=ReadinessDetailsResponse,
)
@step("Call /health/readiness/details without credentials")
def readiness_details_unauthenticated(self) -> Result[ReadinessDetailsResponse]:
return self.proxy.transport.get(
"/health/readiness/details",
@ -87,6 +92,7 @@ class OtherClient:
response_type=ReadinessDetailsResponse,
)
@step("Create the {body.user_role} user {body.user_email} through /user/new")
def user_new(self, body: UserNewBody) -> Result[UserNewResponse]:
"""POST /user/new under the master key: seed the litellm user a JWT
`sub` claim resolves to, before that token ever reaches the proxy."""
@ -97,6 +103,7 @@ class OtherClient:
response_type=UserNewResponse,
)
@step("Read the user's keys from /user/info")
def user_info(self, user_id: str) -> Result[UserInfoWithKeysResponse]:
"""GET /user/info under the master key. Only the user's key rows are
modelled: `token` is the stored key hash, never the plaintext key."""
@ -107,6 +114,7 @@ class OtherClient:
response_type=UserInfoWithKeysResponse,
)
@step("List the JWT-to-key mappings from /jwt/key/mapping/list")
def jwt_mapping_list(self) -> Result[JwtKeyMappingListResponse]:
"""GET /jwt/key/mapping/list under the master key."""
return self.proxy.transport.get(
@ -116,6 +124,7 @@ class OtherClient:
response_type=JwtKeyMappingListResponse,
)
@step("Delete the JWT-to-key mapping")
def jwt_mapping_delete(self, mapping_id: str) -> Result[JwtKeyMappingDeleteResponse]:
"""POST /jwt/key/mapping/delete under the master key."""
return self.proxy.transport.post(
@ -125,6 +134,7 @@ class OtherClient:
response_type=JwtKeyMappingDeleteResponse,
)
@step("Send a /chat/completions request to {body.model} as team {team} with the given token")
def chat_as_team(self, token: str, team: str, body: ChatBody) -> Result[ChatResponse]:
"""POST /chat/completions under `token` with `x-litellm-team-id: team`."""
return self.proxy.transport.post(
@ -137,6 +147,7 @@ class OtherClient:
response_type=ChatResponse,
)
@step("List the models from /v1/models with the given token, in the Anthropic shape: {anthropic}")
def list_models_as(self, token: str, *, anthropic: bool = False) -> Result[ModelsListResponse]:
"""GET /v1/models under `token`, in the OpenAI shape or, with `anthropic`, the
Anthropic Models API shape Claude Code reads. Both carry `data[].id`."""
@ -148,6 +159,7 @@ class OtherClient:
response_type=ModelsListResponse,
)
@step("List users from /user/list with the given key")
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

4
tests/e2e/other/owned_jwt_gateway.py
View file

@ -21,6 +21,7 @@ from typing import Final
from e2e_config import INHERITED_ENV_PREFIXES, available_port
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
@ -36,6 +37,7 @@ class OwnedJwtGateway:
_log_path: Path
_child: subprocess.Popen[bytes] | None = field(default=None, init=False, repr=False)
@step("Start the dedicated JWT proxy and wait for /health/liveliness")
def start(self) -> None:
with self._log_path.open("ab") as log:
self._child = subprocess.Popen(
@ -54,12 +56,14 @@ class OwnedJwtGateway:
time.sleep(0.5)
raise AssertionError("owned JWT gateway did not become ready")
@step("Stop the dedicated JWT proxy")
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("Boot a dedicated proxy {name} with its own litellm_jwtauth config")
def owned_jwt_gateway(
idp: Keycloak, directory: Path, cleanup: ExitStack, *, litellm_jwtauth: str, name: str
) -> OwnedJwtGateway:

25
tests/e2e/other/test_health_lifecycle_e2e.py
View file

@ -17,12 +17,19 @@ import pytest
from e2e_config import MASTER_KEY
from e2e_http import UnauthorizedError, unwrap
from other_client import OtherClient
from e2e_metadata import Domain, Route, Subject, meta
pytestmark = pytest.mark.e2e
class TestHealthLifecycle:
@pytest.mark.covers("other.lifecycle.liveness.ping")
@meta(
Subject(
domain=Domain.DEPLOY_OPS,
route=Route.HEALTH,
)
)
def test_liveness_reports_alive_without_auth(self, client: OtherClient) -> None:
probe = client.liveness()
assert probe.status_code == 200, (
@ -34,6 +41,12 @@ class TestHealthLifecycle:
)
@pytest.mark.covers("other.lifecycle.readiness.public_probe")
@meta(
Subject(
domain=Domain.DEPLOY_OPS,
route=Route.HEALTH,
)
)
def test_readiness_is_reachable_without_credentials(self, client: OtherClient) -> None:
readiness = unwrap(client.readiness_public())
assert readiness.status == "healthy", (
@ -41,6 +54,12 @@ class TestHealthLifecycle:
)
@pytest.mark.covers("other.lifecycle.readiness.reports_db_status")
@meta(
Subject(
domain=Domain.DEPLOY_OPS,
route=Route.HEALTH,
)
)
def test_readiness_reports_connected_db(self, client: OtherClient) -> None:
readiness = unwrap(client.readiness_public())
assert readiness.db == "connected", (
@ -49,6 +68,12 @@ class TestHealthLifecycle:
)
@pytest.mark.covers("other.lifecycle.readiness_details.authenticated_diagnostics")
@meta(
Subject(
domain=Domain.DEPLOY_OPS,
route=Route.HEALTH,
)
)
def test_readiness_details_require_auth_and_expose_diagnostics(self, client: OtherClient) -> None:
anonymous = client.readiness_details_unauthenticated()
assert isinstance(anonymous, UnauthorizedError), (

75
tests/e2e/other/test_jwt_auth_e2e.py
View file

@ -15,6 +15,7 @@ from lifecycle import ResourceManager
from models import ChatBody, ChatMessage, TeamNewBody
from other_client import OtherClient
from pydantic import BaseModel
from e2e_metadata import Domain, Mode, Provider, Subject, meta
pytestmark = pytest.mark.e2e
@ -120,6 +121,14 @@ def _corrupt_signature(token: str) -> str:
class TestJwtAuth:
@pytest.mark.covers("other.auth.jwt.valid_token_allows", "other.auth.jwt.spend_attributed_to_claims")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
providers=(Provider.OPENAI,),
models=(CHEAP_OPENAI_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_valid_token_for_an_existing_team_is_accepted_and_attributed(
self, client: OtherClient, identity: Identity
) -> None:
@ -142,6 +151,13 @@ class TestJwtAuth:
)
@pytest.mark.covers("other.auth.jwt.invalid_signature_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
models=(CHEAP_OPENAI_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_tampered_signature_is_rejected(self, client: OtherClient, identity: Identity) -> None:
tampered: Final = _corrupt_signature(client.idp.access_token(identity))
@ -154,6 +170,13 @@ class TestJwtAuth:
)
@pytest.mark.covers("other.auth.jwt.expired_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
models=(CHEAP_OPENAI_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_expired_token_is_rejected(self, client: OtherClient, identity: Identity) -> None:
expiring: Final = client.idp.access_token(identity, client_id=SHORT_LIVED_CLIENT_ID)
delay: Final = _claims(expiring).exp - time.time() + 1
@ -167,6 +190,13 @@ class TestJwtAuth:
assert "expired" in result.body.lower(), f"the 401 must say the token expired, got {result.body[:300]}"
@pytest.mark.covers("other.auth.jwt.wrong_issuer_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
models=(CHEAP_OPENAI_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_signed_token_from_the_wrong_issuer_is_rejected(self, client: OtherClient, identity: Identity) -> None:
token: Final = client.idp.access_token(identity, issuer_host="unexpected-issuer.invalid")
claims: Final = _claims(token)
@ -177,6 +207,13 @@ class TestJwtAuth:
assert "issuer" in result.body.lower(), f"expected issuer validation to reject the token: {result}"
@pytest.mark.covers("other.auth.jwt.wrong_audience_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
models=(CHEAP_OPENAI_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_signed_token_for_another_application_is_rejected(self, client: OtherClient, identity: Identity) -> None:
token: Final = client.idp.access_token(identity, client_id=WRONG_AUDIENCE_CLIENT_ID)
claims: Final = _claims(token)
@ -189,6 +226,13 @@ class TestJwtAuth:
assert "audience" in result.body.lower(), f"expected audience validation to reject the token: {result}"
@pytest.mark.covers("other.auth.jwt.unknown_team_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
models=(CHEAP_OPENAI_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_token_naming_a_team_that_does_not_exist_is_rejected(
self, client: OtherClient, resources: ResourceManager
) -> None:
@ -204,6 +248,14 @@ class TestJwtAuth:
)
@pytest.mark.covers("other.auth.jwt.virtual_key_unaffected")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
providers=(Provider.OPENAI,),
models=(CHEAP_OPENAI_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_plain_virtual_key_still_works_with_jwt_auth_enabled(self, client: OtherClient, scoped_key: str) -> None:
response: Final = unwrap(client.proxy.chat(scoped_key, _ping()))
assert response.choices, f"an sk- key must keep working on a proxy with enable_jwt_auth, got {response}"
@ -230,6 +282,14 @@ def _denial(client: OtherClient, token: str, team: str) -> str:
class TestJwtTeamHeader:
@pytest.mark.covers("other.auth.jwt.team_header_alias_binds_team")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
providers=(Provider.OPENAI,),
models=(CHEAP_OPENAI_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_team_header_with_the_team_alias_binds_the_same_team_as_the_team_id(
self, client: OtherClient, bound_team: BoundTeam
) -> None:
@ -249,6 +309,14 @@ class TestJwtTeamHeader:
@pytest.mark.covers("other.auth.jwt.team_model_alias_listed_and_routes")
@pytest.mark.parametrize("anthropic", [False, True], ids=["openai_shape", "anthropic_shape"])
@meta(
Subject(
domain=Domain.PROXY_AUTH,
providers=(Provider.OPENAI,),
models=(CHEAP_OPENAI_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_team_model_alias_is_listed_by_v1_models_under_the_same_token_that_routes_it(
self, client: OtherClient, aliased_team: AliasedTeam, anthropic: bool
) -> None:
@ -265,6 +333,13 @@ class TestJwtTeamHeader:
assert aliased_team.target in listed, f"the alias target {aliased_team.target!r} must stay listed, got {listed}"
@pytest.mark.covers("other.auth.jwt.team_header_non_member_alias_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
models=(CHEAP_OPENAI_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_team_header_with_the_alias_of_a_team_the_caller_is_not_in_is_rejected_like_an_unknown_value(
self, client: OtherClient, resources: ResourceManager, bound_team: BoundTeam
) -> None:

25
tests/e2e/other/test_jwt_auto_register_e2e.py
View file

@ -22,6 +22,7 @@ from lifecycle import ResourceManager
from models import ChatBody, ChatMessage, JwtKeyMappingRow, KeyGenerateBody, TeamNewBody, UserNewBody
from other_client import OtherClient
from owned_jwt_gateway import MODEL_NAME, OwnedJwtGateway, owned_jwt_gateway
from e2e_metadata import Domain, Mode, Provider, Subject, meta
pytestmark = pytest.mark.e2e
@ -115,6 +116,14 @@ def minting_gateway(idp: Keycloak, tmp_path_factory: pytest.TempPathFactory) ->
@pytest.mark.owned_gateway
class TestJwtAutoRegisterMapExistingKey:
@pytest.mark.covers("other.auth.jwt.auto_register_maps_existing_key")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
providers=(Provider.GEMINI,),
models=(MODEL_NAME,),
mode=Mode.NONSTREAM,
)
)
def test_first_jwt_call_maps_to_the_users_existing_key_and_mints_none(
self, client: OtherClient, idp: Keycloak, resources: ResourceManager, mapping_gateway: OwnedJwtGateway
) -> None:
@ -144,6 +153,14 @@ class TestJwtAutoRegisterMapExistingKey:
)
@pytest.mark.covers("other.auth.jwt.auto_register_mints_when_keyless")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
providers=(Provider.GEMINI,),
models=(MODEL_NAME,),
mode=Mode.NONSTREAM,
)
)
def test_first_jwt_call_mints_a_key_when_the_user_has_none(
self, client: OtherClient, idp: Keycloak, resources: ResourceManager, mapping_gateway: OwnedJwtGateway
) -> None:
@ -160,6 +177,14 @@ class TestJwtAutoRegisterMapExistingKey:
)
@pytest.mark.covers("other.auth.jwt.auto_register_default_mints")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
providers=(Provider.GEMINI,),
models=(MODEL_NAME,),
mode=Mode.NONSTREAM,
)
)
def test_default_behavior_still_mints_when_the_user_already_has_a_key(
self, client: OtherClient, idp: Keycloak, resources: ResourceManager, minting_gateway: OwnedJwtGateway
) -> None:

11
tests/e2e/other/test_master_key_auth_e2e.py
View file

@ -15,12 +15,18 @@ import pytest
from e2e_config import MASTER_KEY, unique_marker
from e2e_http import UnauthorizedError, unwrap
from other_client import OtherClient
from e2e_metadata import Domain, Subject, meta
pytestmark = pytest.mark.e2e
class TestMasterKeyAuth:
@pytest.mark.covers("other.auth.master_key.valid_allows")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
)
)
def test_master_key_authenticates_and_grants_admin_route(self, client: OtherClient) -> None:
listing = unwrap(client.list_users_as(MASTER_KEY))
assert listing.total >= 0, (
@ -29,6 +35,11 @@ class TestMasterKeyAuth:
)
@pytest.mark.covers("other.auth.master_key.invalid_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
)
)
def test_non_matching_master_key_is_denied(self, client: OtherClient) -> None:
bogus = f"sk-{unique_marker()}"
result = client.list_users_as(bogus)

16
tests/e2e/other/test_session_token_e2e.py
View file

@ -20,6 +20,7 @@ from e2e_http import UnauthorizedError, unwrap
from lifecycle import ResourceManager
from models import KeyGenerateBody, KeyLoggingCallback, KeyLoggingCallbackVars, KeyMetadata
from other_client import OtherClient
from e2e_metadata import Domain, Subject, meta
pytestmark = pytest.mark.e2e
@ -47,12 +48,22 @@ def _admin_session_token(expires_at: datetime) -> str:
class TestSessionToken:
@pytest.mark.covers("other.auth.session_token.valid_allows")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
)
)
def test_unexpired_session_token_reaches_admin_route(self, client: OtherClient) -> None:
token: Final = _admin_session_token(datetime.now(timezone.utc) + timedelta(minutes=10))
listing: Final = unwrap(client.list_users_as(token))
assert listing.total >= 0, f"an unexpired admin session token did not reach /user/list: {listing}"
@pytest.mark.covers("other.auth.session_token.expired_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
)
)
def test_expired_session_token_is_denied(self, client: OtherClient) -> None:
token: Final = _admin_session_token(datetime.now(timezone.utc) - timedelta(minutes=1))
result: Final = client.list_users_as(token)
@ -60,6 +71,11 @@ class TestSessionToken:
assert "expired" in result.body.lower(), f"expected the expired-key error, got {result.body[:300]}"
@pytest.mark.covers("other.auth.session_token.encrypted_value_denied")
@meta(
Subject(
domain=Domain.PROXY_AUTH,
)
)
def test_encrypted_stored_value_is_not_a_bearer_token(
self, client: OtherClient, resources: ResourceManager
) -> None:

4
tests/e2e/secret_manager/secret_store_cyberark.py
View file

@ -10,6 +10,7 @@ from urllib.parse import quote
import pytest
import yaml
from e2e_http import ExternalWrite, Headers, send_text_external
from e2e_metadata import step
from pydantic import Field
from secret_store import SecretBackend
@ -94,6 +95,7 @@ class Conjur:
if not result.ok:
pytest.fail(f"Conjur refused to {action}: HTTP {result.status_code} {result.body[:300]}")
@step("Write the secret {name} to CyberArk Conjur")
def write(self, name: str, value: str) -> None:
self._update_root_policy("POST", f"- !variable {_policy_scalar(name)}\n", f"declare {name}")
result: Final = send_text_external("POST", self._secret_url(name), headers=self._headers(), content=value)
@ -101,6 +103,7 @@ class Conjur:
if not result.ok:
pytest.fail(f"Conjur refused to write {name}: HTTP {result.status_code} {result.body[:300]}")
@step("Read the secret {name} from CyberArk Conjur")
def read(self, name: str) -> str | None:
result: Final = send_text_external("GET", self._secret_url(name), headers=self._headers())
self._fail_unless_reached(result, f"read {name}")
@ -110,6 +113,7 @@ class Conjur:
pytest.fail(f"Conjur refused to read {name}: HTTP {result.status_code} {result.body[:300]}")
return result.body
@step("Delete the secret {name} from CyberArk Conjur")
def destroy(self, name: str) -> None:
self._update_root_policy("PATCH", f"- !delete\n record: !variable {_policy_scalar(name)}\n", f"destroy {name}")

4
tests/e2e/secret_manager/secret_store_hashicorp_vault.py
View file

@ -14,6 +14,7 @@ from e2e_http import (
get_external,
post_json_external,
)
from e2e_metadata import step
from pydantic import BaseModel, Field
from secret_store import SecretBackend
@ -68,6 +69,7 @@ class Vault:
def _metadata_url(self, name: str) -> str:
return f"{self.base_url}/v1/{self.mount}/metadata/{name}"
@step("Write the secret {name} to HashiCorp Vault")
def write(self, name: str, value: str) -> None:
write: Final = post_json_external(
self._data_url(name), headers=self._headers(), json=KvWriteBody(data=KvData(key=value))
@ -77,6 +79,7 @@ class Vault:
if not write.ok:
pytest.fail(f"Vault refused to write {name}: HTTP {write.status_code} {write.body[:300]}")
@step("Read the secret {name} from HashiCorp Vault")
def read(self, name: str) -> str | None:
result: Final = get_external(self._data_url(name), headers=self._headers(), response_type=KvReadResponse)
match result:
@ -89,6 +92,7 @@ class Vault:
case _:
return pytest.fail(f"Vault refused to read {name}: {result}")
@step("Delete the secret {name} from HashiCorp Vault")
def destroy(self, name: str) -> None:
write: Final = delete_external(self._metadata_url(name), headers=self._headers())
if not write.ok and write.status_code != 404:

27
tests/e2e/secret_manager/test_secret_manager_e2e.py
View file

@ -13,6 +13,7 @@ from lifecycle import ResourceManager
from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody, LiteLLMParamsBody
from proxy_client import ProxyClient
from secret_store import SecretStore
from e2e_metadata import Domain, Mode, Provider, Subject, meta
pytestmark = [pytest.mark.e2e, pytest.mark.secret_manager]
@ -73,6 +74,14 @@ def _eventually(proxy: ProxyClient, read: Callable[[], str | None], expected: st
class TestSecretManager:
@pytest.mark.covers("other.config.secret_resolution.kms_integration")
@meta(
Subject(
domain=Domain.DEPLOY_OPS,
providers=(Provider.OPENAI,),
models=(BACKEND_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_deployment_key_resolves_from_the_manager(
self, proxy: ProxyClient, resources: ResourceManager, store: SecretStore, scoped_key: str
) -> None:
@ -83,6 +92,14 @@ class TestSecretManager:
assert response.choices, f"the manager-backed deployment answered with no choices: {response}"
@pytest.mark.covers("other.config.secret_resolution.manager_value_used")
@meta(
Subject(
domain=Domain.DEPLOY_OPS,
providers=(Provider.OPENAI,),
models=(BACKEND_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_deployment_uses_the_value_the_manager_holds(
self, proxy: ProxyClient, resources: ResourceManager, store: SecretStore, scoped_key: str
) -> None:
@ -100,6 +117,11 @@ class TestSecretManager:
pytest.fail(f"expected the provider to reject the manager-held key with 401, got {result}")
@pytest.mark.covers("other.config.secret_manager.virtual_key_stored")
@meta(
Subject(
domain=Domain.DEPLOY_OPS,
)
)
def test_generated_key_is_written_to_the_manager(
self, proxy: ProxyClient, resources: ResourceManager, store: SecretStore
) -> None:
@ -113,6 +135,11 @@ class TestSecretManager:
@pytest.mark.requires_capability("deletes_stored_keys")
@pytest.mark.covers("other.config.secret_manager.virtual_key_deleted")
@meta(
Subject(
domain=Domain.DEPLOY_OPS,
)
)
def test_deleted_key_is_removed_from_the_manager(
self, proxy: ProxyClient, resources: ResourceManager, store: SecretStore
) -> None: