diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 38501c8ddf..e8a60e595c 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -601,7 +601,7 @@ class TestHealthAppFactory: from litellm.proxy.proxy_cli import run_server - runner = CliRunner(mix_stderr=False) + runner = CliRunner() # Mock subprocess.run to simulate prisma being available mock_subprocess_run.return_value = MagicMock(returncode=0) diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 60bdb7d12c..d5d20ce0b0 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -71,20 +71,17 @@ def mock_proxy_config(monkeypatch): @pytest.fixture -def mock_auth(monkeypatch): - """Mock the authentication to bypass auth checks""" +def mock_auth(): + """Mock the authentication to bypass auth checks using FastAPI dependency overrides""" + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app async def mock_user_api_key_auth(): return {"user_id": "test_user"} - from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( - user_api_key_auth, - ) - - monkeypatch.setattr( - "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.user_api_key_auth", - mock_user_api_key_auth, - ) + app.dependency_overrides[user_api_key_auth] = mock_user_api_key_auth + yield + app.dependency_overrides.pop(user_api_key_auth, None) class TestProxySettingEndpoints: @@ -700,6 +697,7 @@ class TestProxySettingEndpoints: def test_get_ui_settings_allows_internal_roles(self, monkeypatch, user_role): """Ensure internal users and viewers can fetch UI settings""" from unittest.mock import AsyncMock, MagicMock + from litellm.proxy.ui_crud_endpoints import proxy_setting_endpoints mock_prisma = MagicMock() @@ -742,8 +740,9 @@ class TestProxySettingEndpoints: ): """Test updating UI settings with an allowlisted field""" from unittest.mock import AsyncMock, MagicMock - from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth # Override the FastAPI dependency with a proper mock mock_user_auth = UserAPIKeyAuth( @@ -782,8 +781,9 @@ class TestProxySettingEndpoints: ): """Test non-allowlisted UI settings are ignored on update""" from unittest.mock import AsyncMock, MagicMock - from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth # Override the FastAPI dependency with a proper mock mock_user_auth = UserAPIKeyAuth( @@ -1105,6 +1105,7 @@ class TestProxySettingEndpoints: def test_get_sso_settings_with_role_mappings(self, mock_proxy_config, mock_auth, monkeypatch): """Test getting SSO settings when role_mappings is present in database""" from unittest.mock import AsyncMock, MagicMock + from litellm.proxy._types import LitellmUserRoles # Mock the prisma client with database record containing role_mappings @@ -1153,6 +1154,7 @@ class TestProxySettingEndpoints: """Test that role_mappings is properly stored and retrieved from SSO settings""" import json from unittest.mock import AsyncMock, MagicMock + from litellm.proxy._types import LitellmUserRoles monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") @@ -1232,8 +1234,9 @@ class TestProxySettingEndpoints: """Test the _setup_role_mappings function directly with custom role mapping logic from environment variables""" import asyncio import os - from litellm.proxy.management_endpoints.ui_sso import _setup_role_mappings + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.ui_sso import _setup_role_mappings # Set up environment variables for custom role mappings using valid Python dict format monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_ROLES", "{'proxy_admin': ['custom-admin-group'], 'internal_user': ['custom-user-group'], 'proxy_admin_viewer': ['custom-viewer-group']}") @@ -1264,6 +1267,7 @@ class TestProxySettingEndpoints: """Test the _setup_role_mappings function returns None when no configuration is available""" import asyncio from unittest.mock import AsyncMock, MagicMock + from litellm.proxy.management_endpoints.ui_sso import _setup_role_mappings # Ensure environment variables are not set @@ -1283,6 +1287,7 @@ class TestProxySettingEndpoints: def test_get_sso_settings_with_env_role_mappings(self, mock_proxy_config, mock_auth, monkeypatch): import json from unittest.mock import AsyncMock, MagicMock + from litellm.proxy._types import LitellmUserRoles monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_ROLES", '{"proxy_admin": ["custom-admin-group"], "internal_user": ["custom-user-group"], "proxy_admin_viewer": ["custom-viewer-group"]}')