From 45e89f977a49b5d9064adeaae2824692a6a3e458 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 3 Sep 2024 15:26:42 -0700
Subject: [PATCH] add test for admin only routes

---
 litellm/proxy/auth/route_checks.py        |  2 +-
 litellm/tests/test_key_generate_prisma.py | 64 +++++++++++++++++++++++
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py
index 57276a7ce81..6c48a7f35e8 100644
--- a/litellm/proxy/auth/route_checks.py
+++ b/litellm/proxy/auth/route_checks.py
@@ -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
 
diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py
index afde1239321..cd4def7657f 100644
--- a/litellm/tests/test_key_generate_prisma.py
+++ b/litellm/tests/test_key_generate_prisma.py
@@ -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