From f8f356fac3bdfb8a52b51d4b9aec7a27c3ea5741 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 16 Apr 2026 16:08:47 -0700 Subject: [PATCH] test: add parametrized tests for api_key value handling in credential check --- tests/proxy_unit_tests/test_proxy_utils.py | 58 +++++++++++++++------- 1 file changed, 41 insertions(+), 17 deletions(-) diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 9f5f14457e..1de0ab450d 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -19,7 +19,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import litellm from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth -from litellm.proxy.auth.auth_utils import is_request_body_safe +from litellm.proxy.auth.auth_utils import ( + check_complete_credentials, + is_request_body_safe, +) from litellm.proxy.litellm_pre_call_utils import ( _get_dynamic_logging_metadata, add_litellm_data_to_request, @@ -33,7 +36,9 @@ def mock_request(monkeypatch): mock_request = Mock(spec=Request) mock_request.query_params = {} # Set mock query_params to an empty dictionary mock_request.headers = {"traceparent": "test_traceparent"} - mock_request.state = State() # Real State so _safe_get_request_headers caching works + mock_request.state = ( + State() + ) # Real State so _safe_get_request_headers caching works monkeypatch.setattr( "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request", mock_request ) @@ -465,6 +470,21 @@ def test_is_request_body_safe_model_enabled( assert expect_error == error_raised +@pytest.mark.parametrize( + "api_key_value, expect_complete", + [ + ("sk-real-key", True), + ("", False), + (None, False), + (" ", False), + ], +) +def test_check_complete_credentials_api_key_values(api_key_value, expect_complete): + request_body = {"model": "gpt-3.5-turbo", "api_key": api_key_value} + result = check_complete_credentials(request_body=request_body) + assert result == expect_complete + + def test_reading_openai_org_id_from_headers(): from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup @@ -735,6 +755,7 @@ def test_get_docs_url(env_vars, expected_url): result = _get_docs_url() assert result == expected_url + @pytest.mark.parametrize( "env_vars, expected_url", [ @@ -1516,7 +1537,7 @@ class MockPrismaClientDB: mock_key_data, ): self.db = MockDb(mock_team_data, mock_key_data) - + async def get_data( self, token: Optional[Union[str, list]] = None, @@ -1534,7 +1555,7 @@ class MockPrismaClientDB: ): """Mock get_data method to return user info for admin""" from litellm.proxy._types import LiteLLM_UserTable - + # Return a proper LiteLLM_UserTable object when querying by user_id if user_id: return LiteLLM_UserTable( @@ -2072,7 +2093,7 @@ def test_team_alias_stale_bypass_disabled_by_default(monkeypatch): monkeypatch.delenv("LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS", raising=False) import litellm.proxy.litellm_pre_call_utils as pre_call_utils from litellm.proxy.litellm_pre_call_utils import _update_model_if_team_alias_exists - + # Reset module-level cache to ensure test isolation pre_call_utils._ENABLE_TEAM_STALE_ALIAS_BYPASS = None @@ -2097,7 +2118,7 @@ def test_team_alias_stale_bypass_disabled_by_default(monkeypatch): def test_team_alias_stale_bypass_enabled_by_flag(monkeypatch): import litellm.proxy.litellm_pre_call_utils as pre_call_utils from litellm.proxy.litellm_pre_call_utils import _update_model_if_team_alias_exists - + # Reset module-level cache to ensure test isolation pre_call_utils._ENABLE_TEAM_STALE_ALIAS_BYPASS = None @@ -2394,16 +2415,17 @@ async def test_handle_logging_proxy_only_error_syncs_normalized_call_type( captured_logging_obj["logging_obj"] = logging_obj return logging_obj, data - with patch( - "litellm.proxy.utils.litellm.utils.function_setup", - side_effect=_capture_function_setup, - ), patch.object( - Logging, "async_failure_handler", new=AsyncMock(return_value=None) - ), patch.object( - Logging, "failure_handler", return_value=None - ), patch( - "litellm.proxy.utils.threading.Thread" - ) as mock_thread: + with ( + patch( + "litellm.proxy.utils.litellm.utils.function_setup", + side_effect=_capture_function_setup, + ), + patch.object( + Logging, "async_failure_handler", new=AsyncMock(return_value=None) + ), + patch.object(Logging, "failure_handler", return_value=None), + patch("litellm.proxy.utils.threading.Thread") as mock_thread, + ): mock_thread.return_value.start = Mock() await proxy_logging._handle_logging_proxy_only_error( @@ -2647,7 +2669,9 @@ async def test_handle_logging_proxy_only_error_skips_handlers_for_pass_through() "model": "claude-3-5-sonnet", } - with patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock) as mock_async: + with patch.object( + logging_obj, "async_failure_handler", new_callable=AsyncMock + ) as mock_async: with patch.object(logging_obj, "failure_handler") as mock_sync: await proxy_logging._handle_logging_proxy_only_error( request_data=request_data,