add unit testing for non_proxy_admin_allowed_routes_check

tests/proxy_admin_ui_tests/test_route_check_unit_tests.py

@@ -14,6 +14,7 @@ import io
 import os
 import time
 
+
 # this file is to test litellm/proxy
 
 sys.path.insert(
@@ -27,6 +28,9 @@ import pytest
 from litellm.proxy.auth.route_checks import RouteChecks
 from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth
 
+# Replace the actual hash_token function with our mock
+import litellm.proxy.auth.route_checks
+
 
 # Mock objects and functions
 class MockRequest:
@@ -34,6 +38,13 @@ class MockRequest:
         self.query_params = query_params or {}
 
 
+def mock_hash_token(token):
+    return token
+
+
+litellm.proxy.auth.route_checks.hash_token = mock_hash_token
+
+
 # Test is_llm_api_route
 def test_is_llm_api_route():
     assert RouteChecks.is_llm_api_route("/v1/chat/completions") is True
@@ -93,3 +104,96 @@ def test_route_matches_pattern():
         )
         is False
     )
+
+
+@pytest.fixture
+def route_checks():
+    return RouteChecks()
+
+
+def test_llm_api_route(route_checks):
+    """
+    Internal User is allowed to access all LLM API routes
+    """
+    assert (
+        route_checks.non_proxy_admin_allowed_routes_check(
+            user_obj=None,
+            _user_role=LitellmUserRoles.INTERNAL_USER.value,
+            route="/v1/chat/completions",
+            request=MockRequest(),
+            valid_token=UserAPIKeyAuth(api_key="test_key"),
+            api_key="test_key",
+            request_data={},
+        )
+        is None
+    )
+
+
+def test_key_info_route_allowed(route_checks):
+    """
+    Internal User is allowed to access /key/info route
+    """
+    assert (
+        route_checks.non_proxy_admin_allowed_routes_check(
+            user_obj=None,
+            _user_role=LitellmUserRoles.INTERNAL_USER.value,
+            route="/key/info",
+            request=MockRequest(query_params={"key": "test_key"}),
+            valid_token=UserAPIKeyAuth(api_key="test_key"),
+            api_key="test_key",
+            request_data={},
+        )
+        is None
+    )
+
+
+def test_key_info_route_forbidden(route_checks):
+    """
+    Internal User is not allowed to access /key/info route for a key they're not using in Authenticated API Key
+    """
+    with pytest.raises(HTTPException) as exc_info:
+        route_checks.non_proxy_admin_allowed_routes_check(
+            user_obj=None,
+            _user_role=LitellmUserRoles.INTERNAL_USER.value,
+            route="/key/info",
+            request=MockRequest(query_params={"key": "wrong_key"}),
+            valid_token=UserAPIKeyAuth(api_key="test_key"),
+            api_key="test_key",
+            request_data={},
+        )
+    assert exc_info.value.status_code == 403
+
+
+def test_user_info_route_allowed(route_checks):
+    """
+    Internal User is allowed to access /user/info route for their own user_id
+    """
+    assert (
+        route_checks.non_proxy_admin_allowed_routes_check(
+            user_obj=None,
+            _user_role=LitellmUserRoles.INTERNAL_USER.value,
+            route="/user/info",
+            request=MockRequest(query_params={"user_id": "test_user"}),
+            valid_token=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
+            api_key="test_key",
+            request_data={},
+        )
+        is None
+    )
+
+
+def test_user_info_route_forbidden(route_checks):
+    """
+    Internal User is not allowed to access /user/info route for a different user_id
+    """
+    with pytest.raises(HTTPException) as exc_info:
+        route_checks.non_proxy_admin_allowed_routes_check(
+            user_obj=None,
+            _user_role=LitellmUserRoles.INTERNAL_USER.value,
+            route="/user/info",
+            request=MockRequest(query_params={"user_id": "wrong_user"}),
+            valid_token=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
+            api_key="test_key",
+            request_data={},
+        )
+    assert exc_info.value.status_code == 403