scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

test(e2e): harden JWT fixtures and cover management lifecycles

Browse source
This commit is contained in:
Yuneng Jiang 2026-09-11 16:36:12 -07:00
parent 2085a37b82
commit 9a80bf2ad4
No known key found for this signature in database
11 changed files with 504 additions and 81 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
tests/e2e/CONTRIBUTING.md
View file

@ -27,16 +27,19 @@ The suites run against a live proxy, so bring one up first by running the litell
2. Bring up a Postgres and a Redis for the proxy to use. The repo-root `docker-compose.yml` already defines a Postgres on `5432`; a `docker run -p 6379:6379 redis:7` covers Redis. Point `DATABASE_URL` / `REDIS_HOST` / `REDIS_PORT` at whatever you run. Tests that read Redis directly default to the deployed shape (TLS + cluster mode) whenever `REDIS_HOST` is set, so for a local standalone Redis also set `REDIS_CLUSTER=false` and `REDIS_SSL=false` (plus `REDIS_PASSWORD` when your Redis requires auth)
3. Start the identity provider the `other/` suite authenticates against, then the litellm proxy against your config, and confirm both are live. It is a real Keycloak, running the realm in `tests/e2e/idp_realm.json`, and the proxy trusts it because `JWT_PUBLIC_KEY_URL` points at that realm's JWKS. The proxy caches the JWKS for `public_key_ttl` (600s) and does not refetch on an unknown `kid`, so restart the proxy whenever you recreate the container:
3. Start the identity provider the JWT API tests authenticate against, then the litellm proxy against your config, and confirm both are live. It is a real Keycloak, running the realm in `tests/e2e/idp_realm.json`, and the proxy trusts it because `JWT_PUBLIC_KEY_URL` points at that realm's JWKS. The proxy caches the JWKS for `public_key_ttl` (600s) and does not refetch on an unknown `kid`, so keep its data volume across restarts; restart the proxy if you deliberately replace that volume:
```bash
set -a && source .env && set +a
docker run -d --name litellm-e2e-idp -p 8480:8080 \
-e KC_BOOTSTRAP_ADMIN_USERNAME=admin -e KC_BOOTSTRAP_ADMIN_PASSWORD=admin \
-v "$PWD/tests/e2e/idp_realm.json:/opt/keycloak/data/import/realm.json:ro" \
-v litellm-e2e-idp-data:/opt/keycloak/data \
quay.io/keycloak/keycloak:26.7.3 start-dev --import-realm
curl -fs --retry 30 --retry-delay 2 --retry-all-errors http://127.0.0.1:8480/realms/litellm-e2e/.well-known/openid-configuration
JWT_PUBLIC_KEY_URL=http://127.0.0.1:8480/realms/litellm-e2e/protocol/openid-connect/certs litellm --config <your-e2e-config>.yml --port 4000
export JWT_ISSUER=http://127.0.0.1:8480/realms/litellm-e2e
export JWT_AUDIENCE=litellm-e2e
export JWT_PUBLIC_KEY_URL="$JWT_ISSUER/protocol/openid-connect/certs"
litellm --config <your-e2e-config>.yml --port 4000
curl -fs http://localhost:4000/health/liveliness
```
@ -53,9 +56,20 @@ The suites run against a live proxy, so bring one up first by running the litell
user_id_upsert: true
```
Leave `JWT_AUDIENCE` and `JWT_ISSUER` unset. Keycloak's access tokens carry `aud: account`, its own audience rather than the proxy's, and their `iss` is whatever base URL the token was requested through, so pinning either one only makes sense once the deployment fixes Keycloak's hostname
Set `JWT_ISSUER` to the exact realm URL used by the test runner and `JWT_AUDIENCE=litellm-e2e`. The realm explicitly maps this audience, `sub`, `email`, and `groups`; the proxy fetches real signing keys from its JWKS endpoint. The rejection tests obtain signed tokens with a different audience or issuer and verify the corresponding rejection reason. The issuer test uses a different HTTP Host when requesting a token from the isolated, dynamically named test IdP.
CI runs this suite against a Keycloak deployed beside the ephemeral stack, and that deployment (the JWT config block, `JWT_PUBLIC_KEY_URL`, and the admin credential handed to the run pod) lives in the project-releaser repo, not here. A stack without it fails the JWT tests rather than skipping them
Keycloak's password grant is a test-only provisioning shortcut, not a production login recommendation. The `litellm-e2e-admin` client adds the proxy's admin scope; the normal client does not. Never reuse this permissive realm outside an isolated test stack.
Management tests can use the shared `idp` and `jwt_identity` fixtures. Each test gets a unique Keycloak group/user and a matching proxy user/team. Setup and fallback cleanup use the master key; the operations and read-backs being tested must explicitly use `caller_key=idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)` (or a member token). See `management/test_jwt_management_e2e.py` for create/read/update/clear/delete and tenant-denial examples. A group claim alone is not database team membership: permission tests explicitly add the member and prove an allowed read before asserting the denied write.
Every successful IdP create immediately registers cleanup, including partial setup failures. Cleanup failures emit warnings. Tokens are minted on demand, and the expiration test waits relative to the token's actual `exp` with a bounded clock-drift check. To check first-attempt behavior locally, run both files with `--reruns 0`:
```bash
E2E_KEYCLOAK_ADMIN_USER=admin E2E_KEYCLOAK_ADMIN_PASSWORD=admin \
uv run pytest tests/e2e/other/test_jwt_auth_e2e.py tests/e2e/management/test_jwt_management_e2e.py --reruns 0 -v
```
CI runs this suite against a Keycloak deployed beside the ephemeral stack, and that deployment (the JWT config block, `JWT_PUBLIC_KEY_URL`, and the admin credential handed to the run pod) lives in the project-releaser repo, not here. CI fetches the realm from the test-runner revision even when it reuses a gateway image from another commit. Keycloak stores its realm, keys and users in a separate schema in the build's PostgreSQL, so replacing the IdP pod preserves token validity. Its startup probe waits for the imported realm. Losing the whole ephemeral database invalidates the stack. Keycloak skips imports into an existing realm, so changes to the realm export require a fresh stack (or deliberately replacing the local data volume). A stack without it fails the JWT tests rather than skipping them
4. Run a suite against it; the harness reads `LITELLM_PROXY_URL` (default `http://localhost:4000`):

37
tests/e2e/conftest.py
View file

@ -17,23 +17,52 @@ import functools
import os
from collections.abc import Generator, Iterator
from datetime import datetime, timezone
from typing import Final
import pytest
import requests
from e2e_config import CONTROL_PLANE_BASE_URL, FIXTURE_DIR, FIXTURE_MODE_RAW, PROXY_BASE_URL
from e2e_config import CONTROL_PLANE_BASE_URL, FIXTURE_DIR, FIXTURE_MODE_RAW, PROXY_BASE_URL, unique_marker
from e2e_db import RESET_OPT_IN_ENV, reset_spend_logs, run_spend_log_cleanup
from e2e_http import unwrap
from fixture_mode import fixture_mode_collection_error, fixture_report_lines
from provider_edge import replay_leftover_error
from idp import Identity, Keycloak, keycloak_from_env
from junit_properties import attach_result_properties
from lifecycle import ProxyClientProvider, ResourceManager
from models import TeamNewBody, UserNewBody, UserNewResponse
from provider_edge import replay_leftover_error
from proxy_client import ProxyClient, build_proxy_client
_E2E_TEST_RAN = pytest.StashKey[bool]()
_CALL_PASSED = pytest.StashKey[bool]()
@pytest.fixture(scope="session")
def idp() -> Keycloak:
return keycloak_from_env()
@pytest.fixture
def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient) -> Identity:
marker: Final = unique_marker()
identity: Final = idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer)
resources.defer(lambda: proxy.delete_user(identity.user_id))
# Seed the canonical user before any JWT call populates the auth cache.
# Group claims grant team access; management membership is added by the test.
unwrap(
proxy.transport.post(
"/user/new",
headers=proxy.transport.master,
json=UserNewBody(
user_id=identity.user_id, user_email=f"{identity.username}@example.com", user_role="internal_user"
),
response_type=UserNewResponse,
)
)
team_id: Final = proxy.create_team(TeamNewBody(team_alias=f"e2e-jwt-{marker}", team_id=identity.group))
resources.defer(lambda: proxy.delete_team(team_id))
return identity
def pytest_configure(config: pytest.Config) -> None:
config.addinivalue_line(
"markers",

4
tests/e2e/coverage_registry/mgmt.yaml
View file

@ -74,3 +74,7 @@
- {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"}
- {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"}
- {id: mgmt.model.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven}
- {id: mgmt.key.jwt.lifecycle, module: mgmt, tier: P0, surface: api, assertions: [lifecycle], source: "management_endpoints/key_management_endpoints.py", rationale: "An IdP-issued admin JWT creates, reads, updates, clears and deletes a key; omitted fields survive updates"}
- {id: mgmt.key.jwt.member_denied, module: mgmt, tier: P0, surface: api, assertions: [member_denied], source: "auth/handle_jwt.py", rationale: "A valid member JWT cannot update an admin-managed key and denial leaves it unchanged"}
- {id: mgmt.key.jwt.other_team_denied, module: mgmt, tier: P0, surface: api, assertions: [other_team_denied], source: "auth/handle_jwt.py", rationale: "A valid JWT for another existing team cannot read the key"}

3
tests/e2e/coverage_registry/other.yaml
View file

@ -51,3 +51,6 @@
- {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"}
- {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"}
- {id: other.a2a.version.serves_pinned_1_0, module: other, tier: P1, area: a2a, assertions: [serves_pinned_1_0], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 1.0 returns the nested 1.0 message shape (result.message with ROLE_AGENT)"}
- {id: other.auth.jwt.wrong_issuer_denied, module: other, tier: P0, area: auth, assertions: [wrong_issuer_denied], source: "auth/handle_jwt.py", rationale: "A signed token with the correct audience and an unexpected issuer is rejected"}
- {id: other.auth.jwt.wrong_audience_denied, module: other, tier: P0, area: auth, assertions: [wrong_audience_denied], source: "auth/handle_jwt.py", rationale: "A signed token from the trusted issuer intended for another app is rejected"}

2
tests/e2e/e2e_http.py
View file

@ -406,6 +406,7 @@ def post_form_external[R: BaseModel](
*,
form: BaseModel,
response_type: type[R],
headers: BaseModel | None = None,
timeout: float = 30.0,
) -> Result[R]:
"""POST an absolute URL outside the proxy as `application/x-www-form-urlencoded`,
@ -415,6 +416,7 @@ def post_form_external[R: BaseModel](
resp = requests.post(
url,
data=_flat(form),
headers=_headers(headers) if headers is not None else None,
timeout=timeout,
)
except requests.RequestException as exc:

78
tests/e2e/idp.py
View file

@ -1,40 +1,15 @@
"""The identity provider the JWT suite authenticates against: a real Keycloak
realm, imported from `idp_realm.json`.
A real IdP rather than a hand-rolled signer because every JWT bug this suite
exists to catch lives in the shape of what an IdP actually emits: `sub` is an
opaque uuid and not a friendly name, group membership arrives as a claim built
by a protocol mapper, the JWKS carries a signing key next to an encryption key
so the proxy has to select on `kid`, and `aud` is the IdP's own audience rather
than the proxy's. A stand-in issuer that mints exactly the claims the tests
assert on can only prove the proxy agrees with the tests.
Tests never hold a signing key. They provision an identity through Keycloak's
admin API (a group named after the litellm team, a user in it with a password
generated for that test alone), then ask Keycloak for an access token through
the direct-access grant, the same way a CLI or service account signs in. The
proxy's `JWT_PUBLIC_KEY_URL` points at this realm's JWKS, so the token the
tests carry is trusted for exactly one reason: Keycloak signed it.
The realm declares two clients. `litellm-e2e-tests` mints ordinary tokens; the
`litellm-e2e-shortlived` client sets `access.token.lifespan` to one second, so
the expiry test lets a genuine token expire instead of forging a stale `exp`.
Connection details come from the environment (`E2E_KEYCLOAK_URL` and the admin
credential). A missing or unreachable IdP is a hard failure naming the start
command, never a skip, so a stack deployed without it turns the run red.
"""
"""Provision isolated identities and obtain signed tokens from the test Keycloak realm."""
from __future__ import annotations
import os
import secrets
from dataclasses import dataclass
import warnings
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Final, Literal
import pytest
from pydantic import BaseModel, Field
from e2e_http import (
AuthHeaders,
ExternalWrite,
@ -45,6 +20,7 @@ from e2e_http import (
post_form_external,
post_json_external,
)
from pydantic import BaseModel, Field
KEYCLOAK_URL_ENV: Final = "E2E_KEYCLOAK_URL"
KEYCLOAK_REALM_ENV: Final = "E2E_KEYCLOAK_REALM"
@ -55,7 +31,8 @@ DEFAULT_KEYCLOAK_URL: Final = "http://127.0.0.1:8480"
DEFAULT_REALM: Final = "litellm-e2e"
TESTS_CLIENT_ID: Final = "litellm-e2e-tests"
SHORT_LIVED_CLIENT_ID: Final = "litellm-e2e-shortlived"
SHORT_LIVED_TOKEN_SECONDS: Final = 1
ADMIN_CLIENT_ID: Final = "litellm-e2e-admin"
WRONG_AUDIENCE_CLIENT_ID: Final = "litellm-e2e-other-app"
_START_HINT: Final = (
"Start it with the `docker run ... quay.io/keycloak/keycloak` command in tests/e2e/CONTRIBUTING.md, "
@ -73,7 +50,11 @@ class TokenGrantForm(BaseModel):
class TokenResponse(BaseModel):
access_token: str
access_token: str = Field(repr=False)
class TokenRequestHeaders(BaseModel):
host: str | None = None
class GroupCreateBody(BaseModel):
@ -105,8 +86,10 @@ class UserCreateBody(BaseModel):
def created_id(write: ExternalWrite, context: str) -> str:
"""The new resource's id, which Keycloak returns only as the last segment of
the Location header on a 201."""
if not write.ok:
if write.status_code != 201:
pytest.fail(f"Keycloak refused to create {context}: HTTP {write.status_code} {write.body[:300]}")
if not write.location or write.location.endswith("/"):
pytest.fail(f"Keycloak created {context} without a resource id in its Location header")
return write.location.rsplit("/", 1)[-1]
@ -117,7 +100,7 @@ class Identity:
user_id: str
username: str
password: str
password: str = field(repr=False)
group: str
group_id: str
@ -127,7 +110,7 @@ class Keycloak:
base_url: str
realm: str
admin_username: str
admin_password: str
admin_password: str = field(repr=False)
@property
def issuer(self) -> str:
@ -183,29 +166,48 @@ class Keycloak:
)
def delete_user(self, user_id: str) -> None:
delete_external(self._admin_url(f"/users/{user_id}"), headers=self._admin_headers())
self._delete(f"/users/{user_id}")
def delete_group(self, group_id: str) -> None:
delete_external(self._admin_url(f"/groups/{group_id}"), headers=self._admin_headers())
self._delete(f"/groups/{group_id}")
def provision(self, *, marker: str, group: str) -> Identity:
def _delete(self, path: str) -> None:
try:
headers: Final = self._admin_headers()
except pytest.fail.Exception as exc:
warnings.warn(f"Keycloak cleanup could not authenticate for {path}: {exc}", RuntimeWarning, stacklevel=2)
return
result: Final = delete_external(self._admin_url(path), headers=headers)
if result.status_code not in (204, 404):
warnings.warn(
f"Keycloak cleanup failed for {path}: HTTP {result.status_code} {result.body[:300]}",
RuntimeWarning,
stacklevel=2,
)
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."""
group_id: Final = self.create_group(group)
defer(lambda: self.delete_group(group_id))
username: Final = f"e2e-jwt-user-{marker}"
password: Final = secrets.token_urlsafe(24)
user_id: Final = self.create_user(
username=username, email=f"{username}@example.com", password=password, group=group
)
defer(lambda: self.delete_user(user_id))
return Identity(user_id=user_id, username=username, password=password, group=group, group_id=group_id)
def access_token(self, identity: Identity, *, client_id: str = TESTS_CLIENT_ID) -> str:
def access_token(
self, identity: Identity, *, client_id: str = TESTS_CLIENT_ID, issuer_host: str | None = None
) -> str:
"""Sign `identity` in through the direct-access grant and hand back the
access token Keycloak signed, exactly as it came off the wire."""
result: Final = post_form_external(
self.token_url(self.realm),
form=TokenGrantForm(client_id=client_id, username=identity.username, password=identity.password),
response_type=TokenResponse,
headers=TokenRequestHeaders(host=issuer_host),
)
return self._token(result, f"a token for {identity.username}")

154
tests/e2e/idp_realm.json
View file

@ -24,7 +24,22 @@
"id.token.claim": "true",
"userinfo.token.claim": "true"
}
},
{
"name": "litellm-audience",
"protocol": "openid-connect",
"protocolMapper": "oidc-audience-mapper",
"consentRequired": false,
"config": {
"included.custom.audience": "litellm-e2e",
"access.token.claim": "true",
"id.token.claim": "false"
}
}
],
"defaultClientScopes": [
"email",
"basic"
]
},
{
@ -49,6 +64,145 @@
"id.token.claim": "true",
"userinfo.token.claim": "true"
}
},
{
"name": "litellm-audience",
"protocol": "openid-connect",
"protocolMapper": "oidc-audience-mapper",
"consentRequired": false,
"config": {
"included.custom.audience": "litellm-e2e",
"access.token.claim": "true",
"id.token.claim": "false"
}
}
],
"defaultClientScopes": [
"email",
"basic"
]
},
{
"clientId": "litellm-e2e-admin",
"enabled": true,
"publicClient": true,
"standardFlowEnabled": false,
"directAccessGrantsEnabled": true,
"defaultClientScopes": [
"email",
"litellm_proxy_admin",
"basic"
],
"protocolMappers": [
{
"name": "groups",
"protocol": "openid-connect",
"protocolMapper": "oidc-group-membership-mapper",
"consentRequired": false,
"config": {
"claim.name": "groups",
"full.path": "false",
"access.token.claim": "true",
"id.token.claim": "true",
"userinfo.token.claim": "true"
}
},
{
"name": "litellm-audience",
"protocol": "openid-connect",
"protocolMapper": "oidc-audience-mapper",
"consentRequired": false,
"config": {
"included.custom.audience": "litellm-e2e",
"access.token.claim": "true",
"id.token.claim": "false"
}
}
]
},
{
"clientId": "litellm-e2e-other-app",
"enabled": true,
"publicClient": true,
"standardFlowEnabled": false,
"directAccessGrantsEnabled": true,
"defaultClientScopes": [
"email",
"basic"
],
"protocolMappers": [
{
"name": "groups",
"protocol": "openid-connect",
"protocolMapper": "oidc-group-membership-mapper",
"consentRequired": false,
"config": {
"claim.name": "groups",
"full.path": "false",
"access.token.claim": "true",
"id.token.claim": "true",
"userinfo.token.claim": "true"
}
},
{
"name": "litellm-audience",
"protocol": "openid-connect",
"protocolMapper": "oidc-audience-mapper",
"consentRequired": false,
"config": {
"included.custom.audience": "litellm-e2e-other-app",
"access.token.claim": "true",
"id.token.claim": "false"
}
}
]
}
],
"clientScopes": [
{
"name": "litellm_proxy_admin",
"protocol": "openid-connect",
"attributes": {
"include.in.token.scope": "true",
"display.on.consent.screen": "false"
}
},
{
"name": "email",
"protocol": "openid-connect",
"attributes": {
"include.in.token.scope": "true"
},
"protocolMappers": [
{
"name": "email",
"protocol": "openid-connect",
"protocolMapper": "oidc-usermodel-property-mapper",
"config": {
"user.attribute": "email",
"claim.name": "email",
"jsonType.label": "String",
"access.token.claim": "true",
"id.token.claim": "true",
"userinfo.token.claim": "true"
}
}
]
},
{
"name": "basic",
"protocol": "openid-connect",
"attributes": {
"include.in.token.scope": "false"
},
"protocolMappers": [
{
"name": "sub",
"protocol": "openid-connect",
"protocolMapper": "oidc-sub-mapper",
"config": {
"access.token.claim": "true"
}
}
]
}

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

@ -10,9 +10,7 @@ import time
from dataclasses import dataclass
import jwt
from e2e_config import MASTER_KEY
from proxy_client import ProxyClient
from e2e_http import (
AuthHeaders,
NetworkError,
@ -37,6 +35,8 @@ from models import (
KeyDeleteBody,
KeyGenerateBody,
KeyGenerateResponse,
KeyInfoParams,
KeyInfoResponse,
KeyListParams,
KeyListResponse,
KeyRegenerateBody,
@ -78,6 +78,7 @@ from models import (
UserNewResponse,
UserUpdateBody,
)
from proxy_client import ProxyClient
MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied"
ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route"
@ -149,13 +150,21 @@ class ManagementClient:
def update_key_models(self, key: str, models: list[str]) -> None:
_ = unwrap(self.update_key(KeyUpdateBody(key=key, models=models)))
def delete_key_strict(self, key: str) -> None:
def key_info_as(self, key: str, *, caller_key: str) -> Result[KeyInfoResponse]:
return self.proxy.transport.get(
"/key/info",
headers=self.proxy.transport.bearer(caller_key),
params=KeyInfoParams(key=key),
response_type=KeyInfoResponse,
)
def delete_key_strict(self, key: str, *, caller_key: str | None = None) -> 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."""
_ = unwrap(
self.proxy.transport.post(
"/key/delete",
headers=self.proxy.transport.master,
headers=self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key),
json=KeyDeleteBody(keys=[key]),
response_type=NoBody,
)

91
tests/e2e/management/test_jwt_management_e2e.py Normal file
View file

@ -0,0 +1,91 @@
"""Management writes and tenant isolation under credentials issued by Keycloak."""
from __future__ import annotations
from typing import Final
import pytest
from e2e_config import CHEAP_OPENAI_MODEL, unique_marker
from e2e_http import UnauthorizedError, UnknownApiError, unwrap
from idp import ADMIN_CLIENT_ID, Identity, Keycloak
from lifecycle import ResourceManager
from management_client import ManagementClient
from models import KeyGenerateBody, KeyUpdateBody, TeamNewBody, UserNewBody
pytestmark = pytest.mark.e2e
class TestJwtManagement:
@pytest.mark.covers("mgmt.key.jwt.lifecycle")
def test_admin_creates_reads_updates_clears_and_deletes_a_key(
self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager
) -> None:
admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)
alias: Final = f"e2e-jwt-key-{unique_marker()}"
created: Final = unwrap(
client.generate_key(
KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group, models=[CHEAP_OPENAI_MODEL]),
caller_key=admin,
)
)
resources.defer(lambda: client.proxy.delete_key(created.key))
original: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info
assert original.key_alias == alias and original.team_id == jwt_identity.group
assert original.models == [CHEAP_OPENAI_MODEL]
updated_alias: Final = f"{alias}-updated"
unwrap(
client.update_key(KeyUpdateBody(key=created.key, key_alias=updated_alias, rpm_limit=120), caller_key=admin)
)
updated: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info
assert updated.key_alias == updated_alias and updated.rpm_limit == 120
assert updated.models == [CHEAP_OPENAI_MODEL], "omitted models must preserve the restriction"
unwrap(client.update_key(KeyUpdateBody(key=created.key, models=[]), caller_key=admin))
cleared: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info
assert cleared.models == [] and cleared.rpm_limit == 120
assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 1
client.delete_key_strict(created.key, caller_key=admin)
assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 0
@pytest.mark.covers("mgmt.key.jwt.member_denied", "mgmt.key.jwt.other_team_denied")
def test_member_cannot_write_and_another_team_cannot_read_the_key(
self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager
) -> None:
admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)
member: Final = idp.access_token(jwt_identity)
alias: Final = f"e2e-jwt-owned-{unique_marker()}"
created: Final = unwrap(
client.generate_key(KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group), caller_key=admin)
)
resources.defer(lambda: client.proxy.delete_key(created.key))
client.add_team_member(jwt_identity.group, jwt_identity.user_id)
assert unwrap(client.key_info_as(created.key, caller_key=member)).info.key_alias == alias
refused: Final = client.update_key(KeyUpdateBody(key=created.key, key_alias="forbidden"), caller_key=member)
assert isinstance(refused, UnauthorizedError), f"member write was accepted: {refused}"
assert "does not have permissions for endpoint" in refused.body.lower(), (
f"expected a permission denial: {refused}"
)
assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.key_alias == alias
marker: Final = unique_marker()
outsider: Final = idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer)
resources.defer(lambda: client.proxy.delete_user(outsider.user_id))
client.create_user(
UserNewBody(
user_id=outsider.user_id, user_email=f"{outsider.username}@example.com", user_role="internal_user"
)
)
team_id: Final = client.proxy.create_team(TeamNewBody(team_alias=marker, team_id=outsider.group))
resources.defer(lambda: client.proxy.delete_team(team_id))
client.add_team_member(outsider.group, outsider.user_id)
outsider_token: Final = idp.access_token(outsider)
hidden: Final = client.key_info_as(created.key, caller_key=outsider_token)
assert isinstance(hidden, UnknownApiError) and hidden.status_code == 403, (
f"another team must not read this key: {hidden}"
)
assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.team_id == jwt_identity.group

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

@ -1,50 +1,42 @@
"""Live e2e: access tokens issued by a real Keycloak realm (idp.py) against a
proxy running with `enable_jwt_auth: true` and the `litellm_jwtauth` block from
CONTRIBUTING.md (sub -> user_id, email -> user_email, groups -> team ids,
user_id_upsert).
Every identity is provisioned in Keycloak for the test that uses it: a group
named after the litellm team, and a user in that group whose password exists
only for the length of the test. Tokens then come from Keycloak's direct-access
grant, so no test ever holds a signing key and the claims the proxy reads are
the ones an IdP really emits (`sub` is Keycloak's user uuid, `groups` comes off
a protocol mapper, `aud` is Keycloak's own audience).
The rejection cases stay honest about where the rejection has to come from: the
bad-signature case corrupts a genuine signature, and the expiry case takes its
token from the realm's one-second client and waits for it to lapse rather than
forging a stale `exp`. Those identities get their own freshly created team, so a
rejection can only be blamed on the token, while the unknown-team case names a
group no litellm team was ever created for. An acceptance is proven twice, at
the boundary (200 from a real provider) and in the spend log the proxy
attributes to the claims. The last case keeps a plain `sk-` virtual key working
on the same proxy, guarding against the flag turning JWT on for everyone.
"""
"""Real Keycloak tokens exercise verification, attribution and virtual-key coexistence."""
from __future__ import annotations
import base64
import time
from typing import Final
import pytest
from e2e_config import CHEAP_OPENAI_MODEL, unique_marker
from e2e_http import UnauthorizedError, UnknownApiError, unwrap
from idp import SHORT_LIVED_CLIENT_ID, SHORT_LIVED_TOKEN_SECONDS, Identity
from idp import SHORT_LIVED_CLIENT_ID, WRONG_AUDIENCE_CLIENT_ID, Identity
from lifecycle import ResourceManager
from models import ChatBody, ChatMessage, TeamNewBody
from other_client import OtherClient
from pydantic import BaseModel
pytestmark = pytest.mark.e2e
class IssuedClaims(BaseModel):
"""Read the IdP's signed payload only to check the test precondition."""
exp: int
sub: str
iss: str
aud: str | list[str]
def _claims(token: str) -> IssuedClaims:
payload: Final = token.split(".")[1]
return IssuedClaims.model_validate_json(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4)))
def _provision(client: OtherClient, resources: ResourceManager, *, marker: str) -> Identity:
"""A Keycloak group and a user in it, torn down with the test. The group name
is what the token's `groups` claim carries, which is what the proxy resolves
as a litellm team id."""
identity: Final = client.idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}")
resources.defer(lambda: client.idp.delete_user(identity.user_id))
resources.defer(lambda: client.idp.delete_group(identity.group_id))
identity: Final = client.idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer)
resources.defer(lambda: client.proxy.delete_user(identity.user_id))
return identity
@ -81,6 +73,7 @@ class TestJwtAuth:
) -> None:
token: Final = client.idp.access_token(identity)
assert _claims(token).sub == identity.user_id, "IdP must emit the provisioned user as sub"
response: Final = unwrap(client.proxy.chat(token, _ping()))
assert response.id is not None and response.choices, (
f"chat under a valid JWT returned no completion: {response}"
@ -111,7 +104,9 @@ class TestJwtAuth:
@pytest.mark.covers("other.auth.jwt.expired_denied")
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)
time.sleep(SHORT_LIVED_TOKEN_SECONDS + 1)
delay: Final = _claims(expiring).exp - time.time() + 1
assert delay <= 5, f"short-lived client expiry or IdP clock drifted: wait would be {delay}s"
time.sleep(max(0, delay))
result: Final = client.proxy.chat(expiring, _ping())
assert isinstance(result, UnauthorizedError), (
@ -119,6 +114,28 @@ 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")
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)
assert claims.iss != client.idp.issuer and "litellm-e2e" in claims.aud
result: Final = client.proxy.chat(token, _ping())
assert isinstance(result, UnauthorizedError), f"wrong issuer must be rejected: {result}"
assert "issuer" in result.body.lower(), f"expected issuer validation to reject the token: {result}"
@pytest.mark.covers("other.auth.jwt.wrong_audience_denied")
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)
assert claims.iss == client.idp.issuer and "litellm-e2e" not in (
[claims.aud] if isinstance(claims.aud, str) else claims.aud
)
result: Final = client.proxy.chat(token, _ping())
assert isinstance(result, UnauthorizedError), f"wrong audience must be rejected: {result}"
assert "audience" in result.body.lower(), f"expected audience validation to reject the token: {result}"
@pytest.mark.covers("other.auth.jwt.unknown_team_denied")
def test_token_naming_a_team_that_does_not_exist_is_rejected(
self, client: OtherClient, resources: ResourceManager

102
tests/e2e/test_idp.py
View file

@ -4,10 +4,14 @@ these carry no `e2e` marker and run everywhere."""
from __future__ import annotations
from collections.abc import Callable, Generator
from contextlib import ExitStack, contextmanager
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from queue import SimpleQueue
from threading import Thread
from typing import Final
import pytest
from e2e_http import ExternalWrite
from idp import (
KEYCLOAK_ADMIN_PASSWORD_ENV,
@ -16,8 +20,8 @@ from idp import (
KEYCLOAK_URL_ENV,
Keycloak,
PasswordCredential,
created_id,
UserCreateBody,
created_id,
keycloak_from_env,
)
@ -44,6 +48,100 @@ def test_a_refused_create_fails_the_test_with_the_idps_own_words() -> None:
created_id(ExternalWrite(status_code=409, body="Group already exists"), "a group")
@pytest.mark.parametrize("location", ["", "http://keycloak/groups/"])
def test_create_without_a_resource_id_fails(location: str) -> None:
with pytest.raises(pytest.fail.Exception, match="resource id"):
created_id(ExternalWrite(status_code=201, location=location), "a group")
@contextmanager
def _idp_server(
*, user_status: int = 201, delete_status: int = 204, admin_status: int = 200
) -> Generator[tuple[Keycloak, SimpleQueue[str]]]:
"""Exercise provisioning failures through the same HTTP transport as live tests."""
deletions: SimpleQueue[str] = SimpleQueue()
class Handler(BaseHTTPRequestHandler):
def log_message(self, format: str, *args: object) -> None:
pass
def do_POST(self) -> None:
self.rfile.read(int(self.headers.get("Content-Length", "0")))
if self.path.endswith("/token"):
self.send_response(admin_status)
self.end_headers()
self.wfile.write(b'{"access_token":"synthetic-harness-token"}')
else:
self.send_response(user_status if self.path.endswith("/users") else 201)
self.send_header("Location", f"{self.path}/resource-1")
self.end_headers()
if user_status != 201 and self.path.endswith("/users"):
self.wfile.write(b"injected create failure")
def do_DELETE(self) -> None:
deletions.put(self.path)
self.send_response(delete_status)
self.end_headers()
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
thread: Final = Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
yield (
Keycloak(
base_url=f"http://127.0.0.1:{server.server_port}",
realm="test",
admin_username="admin",
admin_password="pw",
),
deletions,
)
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
def test_partial_provisioning_removes_the_group_when_user_creation_fails() -> None:
with _idp_server(user_status=500) as (idp, deletions):
with ExitStack() as cleanup:
def defer(callback: Callable[[], object]) -> None:
cleanup.callback(callback)
with pytest.raises(pytest.fail.Exception, match="injected create failure"):
idp.provision(marker="partial", group="team", defer=defer)
assert deletions.get_nowait() == "/admin/realms/test/groups/resource-1"
assert deletions.empty()
def test_successful_provisioning_cleans_up_user_before_group() -> None:
with _idp_server() as (idp, deletions):
with ExitStack() as cleanup:
def defer(callback: Callable[[], object]) -> None:
cleanup.callback(callback)
idp.provision(marker="complete", group="team", defer=defer)
assert deletions.get_nowait() == "/admin/realms/test/users/resource-1"
assert deletions.get_nowait() == "/admin/realms/test/groups/resource-1"
assert deletions.empty()
def test_cleanup_failure_is_visible() -> None:
with _idp_server(delete_status=500) as (idp, _):
with pytest.warns(RuntimeWarning, match="cleanup failed.*HTTP 500"):
idp.delete_group("group")
def test_expired_admin_credentials_do_not_abort_remaining_cleanups() -> None:
with _idp_server(admin_status=401) as (idp, _):
with pytest.warns(RuntimeWarning, match="cleanup could not authenticate") as warnings:
idp.delete_user("user")
idp.delete_group("group")
assert len(warnings) == 2
def test_new_users_are_born_fully_set_up() -> None:
"""A user without a profile or with a pending required action authenticates
nowhere: Keycloak answers every grant with "Account is not fully set up"."""