diff --git a/tests/proxy_admin_ui_tests/test_route_check_unit_tests.py b/tests/proxy_admin_ui_tests/test_route_check_unit_tests.py index edea5f3fe..9a4ec8467 100644 --- a/tests/proxy_admin_ui_tests/test_route_check_unit_tests.py +++ b/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