mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
test(proxy): end-to-end auth_v2 flow through the real FastAPI dependency
In-process integration test with TestClient and an in-memory Prisma stand-in, exercising the real request path without a running proxy or provider keys (the allow/deny decision happens at auth time, before any model call): - flag dispatch -> user_api_key_auth -> v2 entry -> authenticator chain -> casbin -> 403/200 - master key resolves to proxy_admin and passes a governed route - a non-admin with no grant is denied; granting a role and assigning it (live policy CRUD) then unlocks the route, proving the cache reset on write - the _require_admin gate rejects a non-admin over the wire - model calls require the `call` permission: denied without a grant, allowed after granting call on gpt-* and assigning it This is the wiring that unit tests can't cover, verified in CI.
This commit is contained in:
parent
3c1b685a43
commit
0553eda4ab
1 changed files with 158 additions and 0 deletions
158
tests/test_litellm/proxy/auth/v2/test_integration_flow.py
Normal file
158
tests/test_litellm/proxy/auth/v2/test_integration_flow.py
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
"""End-to-end auth_v2 flow through the real FastAPI dependency.
|
||||
|
||||
Exercises flag dispatch -> user_api_key_auth -> v2 entry -> authenticator chain ->
|
||||
casbin authorization -> 403/200, plus live policy CRUD, with an in-memory Prisma
|
||||
stand-in. No running proxy or provider keys: the allow/deny decision happens at
|
||||
auth time, before any model call.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from fastapi import Depends, FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.auth.v2 import policy_store
|
||||
from litellm.proxy.auth.v2.management_endpoints import router as auth_v2_router
|
||||
|
||||
MASTER_KEY = "sk-master-1234"
|
||||
|
||||
|
||||
class _Row:
|
||||
def __init__(self, data):
|
||||
self.ptype = data["ptype"]
|
||||
for i in range(6):
|
||||
setattr(self, f"v{i}", data.get(f"v{i}"))
|
||||
|
||||
|
||||
class _CasbinTable:
|
||||
def __init__(self):
|
||||
self.rows = []
|
||||
|
||||
async def create(self, data):
|
||||
self.rows.append(_Row(data))
|
||||
return data
|
||||
|
||||
async def find_many(self):
|
||||
return list(self.rows)
|
||||
|
||||
async def delete_many(self, where):
|
||||
before = len(self.rows)
|
||||
self.rows = [
|
||||
r
|
||||
for r in self.rows
|
||||
if not all(getattr(r, k, None) == v for k, v in where.items())
|
||||
]
|
||||
return before - len(self.rows)
|
||||
|
||||
|
||||
class _MockPrisma:
|
||||
def __init__(self):
|
||||
self.db = type("_DB", (), {"litellm_casbinrule": _CasbinTable()})()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch):
|
||||
monkeypatch.setattr(ps, "general_settings", {"auth_version": "v2"})
|
||||
monkeypatch.setattr(ps, "master_key", MASTER_KEY)
|
||||
monkeypatch.setattr(ps, "prisma_client", _MockPrisma())
|
||||
policy_store.reset_cache()
|
||||
|
||||
# Non-admin virtual keys resolve to a fixed identity (no DB key row needed).
|
||||
async def fake_get_key_object(hashed_token, **kwargs):
|
||||
return UserAPIKeyAuth(user_id="u_test", token=hashed_token)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.auth.auth_checks.get_key_object", fake_get_key_object
|
||||
)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(auth_v2_router)
|
||||
|
||||
@app.post("/model/new")
|
||||
async def create_model(auth: UserAPIKeyAuth = Depends(user_api_key_auth)):
|
||||
return {"user_id": auth.user_id}
|
||||
|
||||
@app.post("/chat/completions")
|
||||
async def chat(auth: UserAPIKeyAuth = Depends(user_api_key_auth)):
|
||||
return {"user_id": auth.user_id}
|
||||
|
||||
yield TestClient(app, raise_server_exceptions=False)
|
||||
policy_store.reset_cache()
|
||||
|
||||
|
||||
def _h(key):
|
||||
return {"Authorization": f"Bearer {key}"}
|
||||
|
||||
|
||||
def test_master_key_is_admin_and_passes_governed_route(client):
|
||||
# Full chain: flag -> v2 entry -> master-key authn -> proxy_admin -> bootstrap
|
||||
# policy allows -> 200.
|
||||
r = client.post("/model/new", headers=_h(MASTER_KEY))
|
||||
assert r.status_code == 200
|
||||
|
||||
|
||||
def test_non_admin_without_policy_is_denied(client):
|
||||
# u_test has no casbin grant for model:write -> 403.
|
||||
r = client.post("/model/new", headers=_h("sk-user-test"))
|
||||
assert r.status_code == 403
|
||||
|
||||
|
||||
def test_granting_a_role_then_assigning_it_unlocks_the_route(client):
|
||||
# 1. Admin grants role:writer write on models, assigns it to u_test.
|
||||
add_perm = client.post(
|
||||
"/auth/v2/policy/permission/add",
|
||||
headers=_h(MASTER_KEY),
|
||||
json={"role": "writer", "resource": "model", "action": "write"},
|
||||
)
|
||||
assert add_perm.status_code == 200
|
||||
add_assign = client.post(
|
||||
"/auth/v2/policy/assignment/add",
|
||||
headers=_h(MASTER_KEY),
|
||||
json={"subject_type": "user", "subject_id": "u_test", "role": "writer"},
|
||||
)
|
||||
assert add_assign.status_code == 200
|
||||
|
||||
# 2. u_test can now write models (cache was reset on the policy write).
|
||||
r = client.post("/model/new", headers=_h("sk-user-test"))
|
||||
assert r.status_code == 200
|
||||
|
||||
|
||||
def test_non_admin_cannot_administer_policies(client):
|
||||
# The _require_admin gate, over the wire.
|
||||
r = client.post(
|
||||
"/auth/v2/policy/permission/add",
|
||||
headers=_h("sk-user-test"),
|
||||
json={"role": "writer", "resource": "model", "action": "write"},
|
||||
)
|
||||
assert r.status_code == 403
|
||||
|
||||
|
||||
def test_model_call_requires_call_permission(client):
|
||||
# Inference is the `call` action on model:<id>. No grant -> 403.
|
||||
denied = client.post(
|
||||
"/chat/completions", headers=_h("sk-user-test"), json={"model": "gpt-4o"}
|
||||
)
|
||||
assert denied.status_code == 403
|
||||
|
||||
# Grant call on gpt-* to u_test, then it is allowed.
|
||||
client.post(
|
||||
"/auth/v2/policy/permission/add",
|
||||
headers=_h(MASTER_KEY),
|
||||
json={
|
||||
"role": "caller",
|
||||
"resource": "model",
|
||||
"action": "call",
|
||||
"resource_id": "gpt-*",
|
||||
},
|
||||
)
|
||||
client.post(
|
||||
"/auth/v2/policy/assignment/add",
|
||||
headers=_h(MASTER_KEY),
|
||||
json={"subject_type": "user", "subject_id": "u_test", "role": "caller"},
|
||||
)
|
||||
allowed = client.post(
|
||||
"/chat/completions", headers=_h("sk-user-test"), json={"model": "gpt-4o"}
|
||||
)
|
||||
assert allowed.status_code == 200
|
||||
Loading…
Add table
Reference in a new issue