add test for admin only routes

This commit is contained in:
Ishaan Jaff 2024-09-03 15:26:42 -07:00
parent 0b63625673
commit 45e89f977a
2 changed files with 65 additions and 1 deletions

View file

@ -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

View file

@ -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