From c3ab8b3f28b3dde091d5a511093792f73cf70f1a Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 17:37:44 +0530 Subject: [PATCH] fix: add integrations tests to avoid future regressions --- .../proxy/auth/test_auth_integration.py | 395 ++++++++++++++++++ .../proxy/auth/test_route_access_contract.py | 331 +++++++++++++++ 2 files changed, 726 insertions(+) create mode 100644 tests/test_litellm/proxy/auth/test_auth_integration.py create mode 100644 tests/test_litellm/proxy/auth/test_route_access_contract.py diff --git a/tests/test_litellm/proxy/auth/test_auth_integration.py b/tests/test_litellm/proxy/auth/test_auth_integration.py new file mode 100644 index 00000000000..ddbba10e83d --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_auth_integration.py @@ -0,0 +1,395 @@ +""" +Auth Integration Tests +======================= + +End-to-end tests that exercise the full auth pipeline (user_api_key_auth) +for different auth types: custom auth, key-based auth. + +These tests mock at the boundary (DB, cache) but let the real auth logic run, +catching regressions where unit-level mocks might miss interaction bugs. + +Background: The PR chain (#22164 -> #22662 -> b44755db) showed that unit tests +on individual functions weren't enough — the bug was in how functions composed. +These integration tests cover the full flow. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException, Request + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy._types import ( + LiteLLM_UserTable, + LitellmUserRoles, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import common_checks + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def _make_mock_request(route: str, method: str = "POST") -> MagicMock: + """Create a mock FastAPI Request targeting a specific route.""" + req = MagicMock(spec=Request) + req.url.path = route + req.method = method + req.query_params = {} + req.headers = {} + return req + + +# --------------------------------------------------------------------------- +# 1. Custom auth: full pipeline integration +# --------------------------------------------------------------------------- + + +class TestCustomAuthIntegration: + """ + Tests that exercise _run_post_custom_auth_checks through + realistic scenarios, verifying the interaction between: + - custom auth callback + - post-auth checks (end_user lookup, team lookup, etc.) + - common_checks gating + """ + + @pytest.mark.asyncio + async def test_custom_route_allowed_by_default(self): + """ + A custom user-defined route (e.g. /ldap/ngs/ready) must pass + through custom auth without being rejected as admin-only. + + This is the exact scenario that PR #22164 broke. + """ + from litellm.proxy.auth.user_api_key_auth import _run_post_custom_auth_checks + + valid_token = UserAPIKeyAuth( + token="sk-custom-key", + user_id="ldap-user", + ) + + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, patch( + "litellm.proxy.proxy_server.general_settings", + {}, + ): + mock_common.return_value = True + result = await _run_post_custom_auth_checks( + valid_token=valid_token, + request=MagicMock(), + request_data={}, + route="/ldap/ngs/ready", + parent_otel_span=None, + ) + + # common_checks must NOT be called (backwards compat) + mock_common.assert_not_called() + # Token should pass through unchanged + assert result.token == "sk-custom-key" + + @pytest.mark.asyncio + async def test_custom_auth_with_opt_in_enforces_budget(self): + """ + When custom_auth_run_common_checks=True, budget enforcement + via common_checks actually runs. + """ + from litellm.proxy.auth.user_api_key_auth import _run_post_custom_auth_checks + + valid_token = UserAPIKeyAuth( + token="sk-custom-key", + user_id="budget-user", + ) + + budget_error = HTTPException( + status_code=400, + detail="Budget exceeded", + ) + + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, patch( + "litellm.proxy.proxy_server.general_settings", + {"custom_auth_run_common_checks": True}, + ): + mock_common.side_effect = budget_error + + with pytest.raises(HTTPException) as exc_info: + await _run_post_custom_auth_checks( + valid_token=valid_token, + request=MagicMock(), + request_data={}, + route="/chat/completions", + parent_otel_span=None, + ) + + assert exc_info.value.status_code == 400 + assert "Budget exceeded" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_custom_auth_without_opt_in_ignores_budget(self): + """ + Without the opt-in flag, even if common_checks would reject + (budget exceeded), custom auth still passes. + """ + from litellm.proxy.auth.user_api_key_auth import _run_post_custom_auth_checks + + valid_token = UserAPIKeyAuth( + token="sk-custom-key", + user_id="budget-user", + ) + + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, patch( + "litellm.proxy.proxy_server.general_settings", + {}, + ): + # Even if common_checks would raise, it should never be called + mock_common.side_effect = HTTPException( + status_code=400, detail="Budget exceeded" + ) + + result = await _run_post_custom_auth_checks( + valid_token=valid_token, + request=MagicMock(), + request_data={}, + route="/chat/completions", + parent_otel_span=None, + ) + + mock_common.assert_not_called() + assert result.token == "sk-custom-key" + + @pytest.mark.asyncio + async def test_custom_auth_with_end_user_id(self): + """ + Custom auth returning an end_user_id should propagate it + to the valid_token, regardless of common_checks flag. + """ + from litellm.proxy.auth.user_api_key_auth import _run_post_custom_auth_checks + + valid_token = UserAPIKeyAuth( + token="sk-custom-key", + end_user_id="eu-123", + ) + + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, patch( + "litellm.proxy.proxy_server.general_settings", + {}, + ): + mock_common.return_value = True + result = await _run_post_custom_auth_checks( + valid_token=valid_token, + request=MagicMock(), + request_data={}, + route="/v1/chat/completions", + parent_otel_span=None, + ) + + assert result.end_user_id == "eu-123" + + @pytest.mark.asyncio + async def test_custom_auth_with_team_id_triggers_team_lookup(self): + """ + When custom auth sets a team_id on the token, the post-auth + checks should attempt team object lookup. + """ + from litellm.proxy.auth.user_api_key_auth import _run_post_custom_auth_checks + + valid_token = UserAPIKeyAuth( + token="sk-custom-key", + team_id="team-abc", + ) + + mock_team = MagicMock() + mock_team.blocked = False + + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=mock_team, + ) as mock_get_team, patch( + "litellm.proxy.proxy_server.general_settings", + {"custom_auth_run_common_checks": True}, + ), patch( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", + MagicMock(), + ): + mock_common.return_value = True + await _run_post_custom_auth_checks( + valid_token=valid_token, + request=MagicMock(), + request_data={}, + route="/chat/completions", + parent_otel_span=None, + ) + + mock_get_team.assert_called_once() + mock_common.assert_called_once() + + +# --------------------------------------------------------------------------- +# 2. Key-based auth: common_checks integration +# --------------------------------------------------------------------------- + + +class TestKeyBasedAuthCommonChecks: + """ + Tests that common_checks correctly enforces route authorization + for standard key-based auth flows. + """ + + @pytest.mark.asyncio + async def test_llm_route_allowed_for_regular_user(self): + """Regular users can call LLM API routes through common_checks.""" + user_obj = LiteLLM_UserTable( + user_id="user1", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + valid_token = UserAPIKeyAuth( + token="sk-key", + user_id="user1", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + result = await common_checks( + request_body={"model": "gpt-4"}, + team_object=None, + user_object=user_obj, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=_make_mock_request("/chat/completions"), + ) + assert result is True + + @pytest.mark.asyncio + async def test_admin_route_denied_for_regular_user(self): + """Regular users cannot call admin-only routes through common_checks.""" + from litellm.proxy.auth.auth_checks import common_checks + + user_obj = LiteLLM_UserTable( + user_id="user1", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + valid_token = UserAPIKeyAuth( + token="sk-key", + user_id="user1", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + with pytest.raises(Exception): + await common_checks( + request_body={}, + team_object=None, + user_object=user_obj, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/config/update", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=_make_mock_request("/config/update"), + ) + + @pytest.mark.asyncio + async def test_admin_can_access_any_route(self): + """Proxy admin can call any route through common_checks.""" + from litellm.proxy.auth.auth_checks import common_checks + + user_obj = LiteLLM_UserTable( + user_id="admin1", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ) + valid_token = UserAPIKeyAuth( + token="sk-admin", + user_id="admin1", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ) + + for route in ["/chat/completions", "/config/update", "/key/generate"]: + result = await common_checks( + request_body={}, + team_object=None, + user_object=user_obj, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route=route, + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=_make_mock_request(route), + ) + assert result is True + + +# --------------------------------------------------------------------------- +# 3. Cross-cutting: ensure custom_auth_run_common_checks flag isolation +# --------------------------------------------------------------------------- + + +class TestFlagIsolation: + """ + Verify that the custom_auth_run_common_checks flag ONLY affects + custom auth flows and does NOT interfere with key-based or JWT auth. + """ + + @pytest.mark.asyncio + async def test_flag_does_not_affect_common_checks_directly(self): + """ + common_checks() itself should not read the flag — + it's the caller (_run_post_custom_auth_checks) that gates the call. + """ + from litellm.proxy.auth.auth_checks import common_checks + + user_obj = LiteLLM_UserTable( + user_id="user1", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + valid_token = UserAPIKeyAuth( + token="sk-key", + user_id="user1", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + # common_checks should work the same regardless of the flag + for flag_value in [True, False]: + result = await common_checks( + request_body={}, + team_object=None, + user_object=user_obj, + end_user_object=None, + global_proxy_spend=None, + general_settings={"custom_auth_run_common_checks": flag_value}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=_make_mock_request("/chat/completions"), + ) + assert result is True diff --git a/tests/test_litellm/proxy/auth/test_route_access_contract.py b/tests/test_litellm/proxy/auth/test_route_access_contract.py new file mode 100644 index 00000000000..138e26b4fbe --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_route_access_contract.py @@ -0,0 +1,331 @@ +""" +Route Access Contract Tests +============================ + +Parameterized matrix test that encodes which auth_type x role x route combinations +are allowed or denied. This acts as a regression safety net — any PR that changes +auth behavior will break specific rows and force the author to explicitly update +the contract. + +Background: +- PR #22164 added common_checks() to custom auth, breaking custom routes. +- PR #22662 patched it with skip_route_check. +- PR (b44755db) replaced both with an opt-in flag. +All three PRs would have been caught (or made unnecessary) by this test. +""" + +import os +import sys +from typing import Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import Request + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy._types import ( + LiteLLM_UserTable, + LitellmUserRoles, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import _is_allowed_route, common_checks +from litellm.proxy.auth.route_checks import RouteChecks + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_request() -> MagicMock: + req = MagicMock(spec=Request) + req.query_params = {} + return req + + +def _make_user(role: str) -> LiteLLM_UserTable: + return LiteLLM_UserTable( + user_id="test_user", + user_email="test@example.com", + user_role=role, + ) + + +def _make_token( + role: Optional[str] = None, + team_id: Optional[str] = None, + allowed_routes: Optional[list] = None, +) -> UserAPIKeyAuth: + kwargs = {"token": "sk-test", "user_id": "test_user"} + if role: + kwargs["user_role"] = role + if team_id: + kwargs["team_id"] = team_id + if allowed_routes: + kwargs["allowed_routes"] = allowed_routes + return UserAPIKeyAuth(**kwargs) + + +# --------------------------------------------------------------------------- +# 1. Route check contract for key-based auth (non-admin roles) +# Tests _is_allowed_route / RouteChecks.non_proxy_admin_allowed_routes_check +# --------------------------------------------------------------------------- + +# (role, route, expected_outcome) +# "allowed" = no exception raised, "denied" = HTTPException / Exception raised +_KEY_BASED_ROUTE_MATRIX = [ + # --- LLM API routes: all non-admin roles can call them (except admin_view_only) --- + (LitellmUserRoles.INTERNAL_USER.value, "/chat/completions", "allowed"), + (LitellmUserRoles.INTERNAL_USER.value, "/v1/chat/completions", "allowed"), + (LitellmUserRoles.INTERNAL_USER.value, "/v1/embeddings", "allowed"), + (LitellmUserRoles.INTERNAL_USER.value, "/v1/models", "allowed"), + (LitellmUserRoles.INTERNAL_USER.value, "/v1/messages", "allowed"), + # --- Info routes: accessible to all non-admin roles --- + (LitellmUserRoles.INTERNAL_USER.value, "/model/info", "allowed"), + (LitellmUserRoles.INTERNAL_USER.value, "/key/info", "allowed"), + (LitellmUserRoles.INTERNAL_USER.value, "/health", "allowed"), + (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, "/model/info", "allowed"), + (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, "/key/info", "allowed"), + # --- Management write routes: denied for non-admin --- + (LitellmUserRoles.INTERNAL_USER.value, "/user/new", "denied"), + (LitellmUserRoles.INTERNAL_USER.value, "/user/delete", "denied"), + (LitellmUserRoles.INTERNAL_USER.value, "/team/new", "denied"), + (LitellmUserRoles.INTERNAL_USER.value, "/team/delete", "denied"), + (LitellmUserRoles.INTERNAL_USER.value, "/config/update", "denied"), + # internal_user CAN access key_management_routes and self_managed_routes + (LitellmUserRoles.INTERNAL_USER.value, "/key/generate", "allowed"), + (LitellmUserRoles.INTERNAL_USER.value, "/model/new", "allowed"), + # Admin view-only: write management routes denied + (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, "/user/new", "denied"), + (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, "/key/generate", "denied"), + (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, "/team/delete", "denied"), + # Admin view-only: read management routes allowed + (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, "/user/list", "allowed"), + (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, "/user/info", "allowed"), + (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, "/team/info", "allowed"), + # --- Internal user: spend tracking routes allowed --- + (LitellmUserRoles.INTERNAL_USER.value, "/global/spend/tags", "allowed"), + (LitellmUserRoles.INTERNAL_USER.value, "/global/spend/keys", "allowed"), + # Internal user view-only: spend routes allowed, management denied + (LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, "/spend/keys", "allowed"), + (LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, "/global/spend/logs", "allowed"), + (LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, "/key/generate", "denied"), + (LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, "/user/new", "denied"), + # --- Self-managed routes: accessible to non-admin --- + (LitellmUserRoles.INTERNAL_USER.value, "/team/member_add", "allowed"), + (LitellmUserRoles.INTERNAL_USER.value, "/team/member_delete", "allowed"), +] + + +@pytest.mark.parametrize( + "role,route,expected", + _KEY_BASED_ROUTE_MATRIX, + ids=[f"{r}|{rt}|{exp}" for r, rt, exp in _KEY_BASED_ROUTE_MATRIX], +) +def test_key_based_route_access(role, route, expected): + """Contract test: key-based auth route access for each role.""" + request = _make_request() + user_obj = _make_user(role) + valid_token = _make_token(role=role) + + if expected == "allowed": + # Should NOT raise + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=role, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + else: + with pytest.raises(Exception): + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=role, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + + +# --------------------------------------------------------------------------- +# 2. Proxy admin bypasses all route checks +# --------------------------------------------------------------------------- + +_ADMIN_ROUTES = [ + "/chat/completions", + "/key/generate", + "/user/new", + "/config/update", + "/global/spend/reset", + "/model/delete", +] + + +@pytest.mark.parametrize("route", _ADMIN_ROUTES) +def test_proxy_admin_can_access_all_routes(route): + """Proxy admin should pass _is_allowed_route for any route.""" + request = _make_request() + user_obj = _make_user(LitellmUserRoles.PROXY_ADMIN.value) + valid_token = _make_token(role=LitellmUserRoles.PROXY_ADMIN.value) + + result = _is_allowed_route( + route=route, + token_type="api", + request=request, + request_data={}, + valid_token=valid_token, + user_obj=user_obj, + ) + assert result is True + + +# --------------------------------------------------------------------------- +# 3. Custom auth: common_checks opt-in contract +# This is the exact regression that caused the 3-PR chain. +# --------------------------------------------------------------------------- + +_CUSTOM_AUTH_ROUTES = [ + "/chat/completions", + "/v1/chat/completions", + "/ldap/ngs/ready", # custom user-defined route (the one that broke) + "/my-app/webhook", # another hypothetical custom route + "/key/info", + "/health", +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", _CUSTOM_AUTH_ROUTES) +async def test_custom_auth_default_skips_common_checks(route): + """ + Default custom auth: common_checks is NOT called, so ANY route passes. + This is the backwards-compatible behavior (pre-#22164). + """ + from litellm.proxy.auth.user_api_key_auth import _run_post_custom_auth_checks + + valid_token = UserAPIKeyAuth(token="sk-custom") + mock_request = MagicMock() + + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, patch( + "litellm.proxy.proxy_server.general_settings", + {}, + ): + mock_common.return_value = True + await _run_post_custom_auth_checks( + valid_token=valid_token, + request=mock_request, + request_data={}, + route=route, + parent_otel_span=None, + ) + mock_common.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", _CUSTOM_AUTH_ROUTES) +async def test_custom_auth_opt_in_runs_common_checks(route): + """ + With custom_auth_run_common_checks=True, common_checks IS called. + """ + from litellm.proxy.auth.user_api_key_auth import _run_post_custom_auth_checks + + valid_token = UserAPIKeyAuth(token="sk-custom") + mock_request = MagicMock() + + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, patch( + "litellm.proxy.proxy_server.general_settings", + {"custom_auth_run_common_checks": True}, + ): + mock_common.return_value = True + await _run_post_custom_auth_checks( + valid_token=valid_token, + request=mock_request, + request_data={}, + route=route, + parent_otel_span=None, + ) + mock_common.assert_called_once() + + +# --------------------------------------------------------------------------- +# 4. common_checks route enforcement (key-based path) +# Verifies that _is_allowed_route is always called inside common_checks +# (i.e. the skip_route_check parameter from #22662 stays removed). +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_common_checks_always_runs_route_check(): + """ + Regression: common_checks must always run route authorization. + The skip_route_check param (PR #22662) must NOT come back. + """ + import inspect + + sig = inspect.signature(common_checks) + param_names = list(sig.parameters.keys()) + + # skip_route_check must NOT exist as a parameter + assert "skip_route_check" not in param_names, ( + "skip_route_check parameter found in common_checks — " + "this was removed in favor of the custom_auth_run_common_checks flag" + ) + + +# --------------------------------------------------------------------------- +# 5. Virtual key allowed_routes override +# --------------------------------------------------------------------------- + +_VIRTUAL_KEY_CASES = [ + # (allowed_routes on key, route being accessed, expected) + (["/chat/completions", "/key/info"], "/chat/completions", "allowed"), + # allowed_routes on the key restricts to only those routes; + # /config/update is not in any non-admin allowed set, so it's denied + (["/chat/completions"], "/config/update", "denied"), + (["llm_api_routes"], "/v1/chat/completions", "allowed"), +] + + +@pytest.mark.parametrize( + "allowed_routes,route,expected", + _VIRTUAL_KEY_CASES, + ids=[f"{ar}|{rt}|{exp}" for ar, rt, exp in _VIRTUAL_KEY_CASES], +) +def test_virtual_key_allowed_routes(allowed_routes, route, expected): + """Keys with explicit allowed_routes restrict access to those routes only.""" + request = _make_request() + user_obj = _make_user(LitellmUserRoles.INTERNAL_USER.value) + valid_token = _make_token( + role=LitellmUserRoles.INTERNAL_USER.value, + allowed_routes=allowed_routes, + ) + + if expected == "allowed": + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + else: + with pytest.raises(Exception): + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + )