mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
add test for admin only routes
This commit is contained in:
parent
0b63625673
commit
45e89f977a
2 changed files with 65 additions and 1 deletions
|
|
@ -133,7 +133,7 @@ def custom_admin_only_route_check(route: str):
|
|||
if route in general_settings["admin_only_routes"]:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="user not allowed to access this route. This is an admin only route",
|
||||
detail=f"user not allowed to access this route. Route={route} is an admin only route",
|
||||
)
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -3092,3 +3092,67 @@ async def test_team_tags(prisma_client):
|
|||
|
||||
print("team_info_response", team_info_response)
|
||||
assert team_info_response["team_info"].metadata["tags"] == ["teamA", "teamB"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_only_routes(prisma_client):
|
||||
"""
|
||||
Tests if setting admin_only_routes works
|
||||
|
||||
only an admin should be able to access admin only routes
|
||||
"""
|
||||
litellm.set_verbose = True
|
||||
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
||||
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||
await litellm.proxy.proxy_server.prisma_client.connect()
|
||||
general_settings = {
|
||||
"allowed_routes": ["/embeddings", "/key/generate"],
|
||||
"admin_only_routes": ["/key/generate"],
|
||||
}
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
setattr(proxy_server, "general_settings", general_settings)
|
||||
|
||||
admin_user = await new_user(
|
||||
data=NewUserRequest(
|
||||
user_name="admin",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
non_admin_user = await new_user(
|
||||
data=NewUserRequest(
|
||||
user_name="non-admin",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
admin_user_key = admin_user.key
|
||||
non_admin_user_key = non_admin_user.key
|
||||
|
||||
assert admin_user_key is not None
|
||||
assert non_admin_user_key is not None
|
||||
|
||||
# assert non-admin can not access admin routes
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/key/generate")
|
||||
await user_api_key_auth(
|
||||
request=request,
|
||||
api_key="Bearer " + admin_user_key,
|
||||
)
|
||||
|
||||
# this should pass
|
||||
|
||||
try:
|
||||
await user_api_key_auth(
|
||||
request=request,
|
||||
api_key="Bearer " + non_admin_user_key,
|
||||
)
|
||||
pytest.fail("Expected this call to fail. User is over limit.")
|
||||
except Exception as e:
|
||||
print("error str=", str(e.message))
|
||||
error_str = str(e.message)
|
||||
assert "Route" in error_str and "admin only route" in error_str
|
||||
pass
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue