From 9a338e1b6b1ee91d15d79e829be15a730c59e700 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 5 May 2026 17:21:18 -0700 Subject: [PATCH] [Test] Tests: Stop parametrizing API keys into pytest test IDs (#27249) Several tests parametrized over (model, api_key, ...) tuples or raw token strings, causing pytest to embed those values in the test ID and print them in CI logs. Refactored each affected test to keep the same coverage without putting key material into parametrize. - audio_tests/test_audio_speech.py: split env-var keys into separate azure/openai test functions sharing a helper; sync_mode parametrize preserved. - audio_tests/test_whisper.py: split into openai_whisper / azure_whisper functions sharing a helper; response_format parametrize preserved. - local_testing/test_embedding.py: single-case parametrize inlined. - proxy_unit_tests/test_user_api_key_auth.py: 5 header parametrize cases split into 5 named tests sharing an _assert helper. - proxy_unit_tests/test_proxy_utils.py: 4 api_key_value cases split into 4 named tests. - test_litellm/proxy/auth/test_user_api_key_auth.py: 5 key-prefix cases (Bearer / Basic / lowercase bearer / raw / AWS SigV4) split into 5 named tests. Verified: black clean; 14 refactored unit tests pass; pytest collects audio/embedding tests with safe IDs (no key material in test IDs). --- tests/audio_tests/test_audio_speech.py | 43 +++++++----- tests/audio_tests/test_whisper.py | 51 +++++++++----- tests/local_testing/test_embedding.py | 15 ++-- tests/proxy_unit_tests/test_proxy_utils.py | 29 +++++--- .../test_user_api_key_auth.py | 69 ++++++++++++------- .../proxy/auth/test_user_api_key_auth.py | 60 +++++++++++----- 6 files changed, 170 insertions(+), 97 deletions(-) diff --git a/tests/audio_tests/test_audio_speech.py b/tests/audio_tests/test_audio_speech.py index 46d4515891..52a2316a16 100644 --- a/tests/audio_tests/test_audio_speech.py +++ b/tests/audio_tests/test_audio_speech.py @@ -26,24 +26,7 @@ import pytest import litellm -@pytest.mark.parametrize( - "sync_mode", - [True, False], -) -@pytest.mark.parametrize( - "model, api_key, api_base", - [ - ( - "azure/tts", - os.getenv("AZURE_TTS_API_KEY"), - os.getenv("AZURE_TTS_API_BASE"), - ), - ("openai/tts-1", os.getenv("OPENAI_API_KEY"), None), - ], -) # , -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=1) -async def test_audio_speech_litellm(sync_mode, model, api_base, api_key): +async def _run_audio_speech_litellm(sync_mode, model, api_base, api_key): litellm._turn_on_debug() speech_file_path = Path(__file__).parent / "speech.mp3" @@ -85,6 +68,30 @@ async def test_audio_speech_litellm(sync_mode, model, api_base, api_key): assert isinstance(response, HttpxBinaryResponseContent) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) +async def test_audio_speech_litellm_azure(sync_mode): + await _run_audio_speech_litellm( + sync_mode=sync_mode, + model="azure/tts", + api_base=os.getenv("AZURE_TTS_API_BASE"), + api_key=os.getenv("AZURE_TTS_API_KEY"), + ) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) +async def test_audio_speech_litellm_openai(sync_mode): + await _run_audio_speech_litellm( + sync_mode=sync_mode, + model="openai/tts-1", + api_base=None, + api_key=os.getenv("OPENAI_API_KEY"), + ) + + @pytest.mark.parametrize( "sync_mode", [False, True], diff --git a/tests/audio_tests/test_whisper.py b/tests/audio_tests/test_whisper.py index 199b0e6a4f..cdf079f8cb 100644 --- a/tests/audio_tests/test_whisper.py +++ b/tests/audio_tests/test_whisper.py @@ -39,24 +39,7 @@ import litellm from litellm import Router -@pytest.mark.parametrize( - "model, api_key, api_base", - [ - ("whisper-1", None, None), - ( - "azure/whisper", - os.getenv("AZURE_WHISPER_API_KEY"), - os.getenv("AZURE_WHISPER_API_BASE"), - ), - ], -) -@pytest.mark.parametrize( - "response_format, timestamp_granularities", - [("json", None), ("vtt", None), ("verbose_json", ["word"])], -) -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=1) -async def test_transcription( +async def _run_transcription( model, api_key, api_base, response_format, timestamp_granularities ): transcript = await litellm.atranscription( @@ -74,6 +57,38 @@ async def test_transcription( assert transcript.text is not None +@pytest.mark.parametrize( + "response_format, timestamp_granularities", + [("json", None), ("vtt", None), ("verbose_json", ["word"])], +) +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) +async def test_transcription_openai_whisper(response_format, timestamp_granularities): + await _run_transcription( + model="whisper-1", + api_key=None, + api_base=None, + response_format=response_format, + timestamp_granularities=timestamp_granularities, + ) + + +@pytest.mark.parametrize( + "response_format, timestamp_granularities", + [("json", None), ("vtt", None), ("verbose_json", ["word"])], +) +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) +async def test_transcription_azure_whisper(response_format, timestamp_granularities): + await _run_transcription( + model="azure/whisper", + api_key=os.getenv("AZURE_WHISPER_API_KEY"), + api_base=os.getenv("AZURE_WHISPER_API_BASE"), + response_format=response_format, + timestamp_granularities=timestamp_granularities, + ) + + @pytest.mark.asyncio() async def test_transcription_caching(): import litellm diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index 82510b6f4f..14626aa8e4 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -193,19 +193,12 @@ def _azure_ai_image_mock_response(*args, **kwargs): return new_response -@pytest.mark.parametrize( - "model, api_base, api_key", - [ - ( - "azure_ai/Cohere-embed-v3-multilingual-2", - os.getenv("AZURE_AI_API_BASE"), - os.getenv("AZURE_AI_API_KEY"), - ) - ], -) @pytest.mark.parametrize("sync_mode", [True]) # , False @pytest.mark.asyncio -async def test_azure_ai_embedding_image(model, api_base, api_key, sync_mode): +async def test_azure_ai_embedding_image(sync_mode): + model = "azure_ai/Cohere-embed-v3-multilingual-2" + api_base = os.getenv("AZURE_AI_API_BASE") + api_key = os.getenv("AZURE_AI_API_KEY") try: os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 3332e77f2d..70232f25c3 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -501,21 +501,30 @@ 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): +def _assert_check_complete_credentials(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_check_complete_credentials_with_real_key(): + _assert_check_complete_credentials( + api_key_value="sk-" + "x" * 8, expect_complete=True + ) + + +def test_check_complete_credentials_with_empty_string(): + _assert_check_complete_credentials(api_key_value="", expect_complete=False) + + +def test_check_complete_credentials_with_none(): + _assert_check_complete_credentials(api_key_value=None, expect_complete=False) + + +def test_check_complete_credentials_with_whitespace(): + _assert_check_complete_credentials(api_key_value=" ", expect_complete=False) + + def test_reading_openai_org_id_from_headers(): from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 543cabb6b4..210347aaf9 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -269,9 +269,7 @@ async def test_aaauser_personal_budgets(key_ownership): test_user_cache = getattr(litellm.proxy.proxy_server, "user_api_key_cache") assert ( - test_user_cache.get_cache( - key=hash_token(user_key), model_type=UserAPIKeyAuth - ) + test_user_cache.get_cache(key=hash_token(user_key), model_type=UserAPIKeyAuth) == valid_token ) @@ -514,36 +512,59 @@ async def test_auth_not_connected_to_db(): assert valid_token.token == "failed-to-connect-to-db" -@pytest.mark.parametrize( - "headers, custom_header_name, expected_api_key", - [ - # Test with valid Bearer token - ({"x-custom-api-key": "Bearer sk-12345678"}, "x-custom-api-key", "sk-12345678"), - # Test with raw token (no Bearer prefix) - ({"x-custom-api-key": "Bearer sk-12345678"}, "x-custom-api-key", "sk-12345678"), - # Test with empty header value - ({"x-custom-api-key": ""}, "x-custom-api-key", ""), - # Test with missing header - ({}, "X-Custom-API-Key", ""), - # Test with different header casing - ({"X-CUSTOM-API-KEY": "Bearer sk-12345678"}, "X-Custom-API-Key", "sk-12345678"), - ], -) -def test_get_api_key_from_custom_header(headers, custom_header_name, expected_api_key): +def _assert_api_key_from_custom_header(headers, custom_header_name, expected_api_key): verbose_proxy_logger.setLevel(logging.DEBUG) - - # Mock the Request object request = MagicMock(spec=Request) request.headers = headers - - # Call the function and verify it doesn't raise an exception - api_key = get_api_key_from_custom_header( request=request, custom_litellm_key_header_name=custom_header_name ) assert api_key == expected_api_key +def test_get_api_key_from_custom_header_bearer_token(): + token = "sk-" + "1" * 8 + _assert_api_key_from_custom_header( + headers={"x-custom-api-key": f"Bearer {token}"}, + custom_header_name="x-custom-api-key", + expected_api_key=token, + ) + + +def test_get_api_key_from_custom_header_raw_token(): + token = "sk-" + "1" * 8 + _assert_api_key_from_custom_header( + headers={"x-custom-api-key": f"Bearer {token}"}, + custom_header_name="x-custom-api-key", + expected_api_key=token, + ) + + +def test_get_api_key_from_custom_header_empty_value(): + _assert_api_key_from_custom_header( + headers={"x-custom-api-key": ""}, + custom_header_name="x-custom-api-key", + expected_api_key="", + ) + + +def test_get_api_key_from_custom_header_missing_header(): + _assert_api_key_from_custom_header( + headers={}, + custom_header_name="X-Custom-API-Key", + expected_api_key="", + ) + + +def test_get_api_key_from_custom_header_different_casing(): + token = "sk-" + "1" * 8 + _assert_api_key_from_custom_header( + headers={"X-CUSTOM-API-KEY": f"Bearer {token}"}, + custom_header_name="X-Custom-API-Key", + expected_api_key=token, + ) + + from litellm.proxy._types import LitellmUserRoles diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 2e5eef2a0a..95b3d746c6 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -556,22 +556,7 @@ async def test_enterprise_custom_auth_runs_post_custom_auth_checks_when_opt_in() litellm.enable_post_custom_auth_checks = original_flag -@pytest.mark.parametrize( - "custom_litellm_key_header, api_key, passed_in_key", - [ - ("Bearer sk-12345678", "sk-12345678", "Bearer sk-12345678"), - ("Basic sk-12345678", "sk-12345678", "Basic sk-12345678"), - ("bearer sk-12345678", "sk-12345678", "bearer sk-12345678"), - ("sk-12345678", "sk-12345678", "sk-12345678"), - # AWS Signature V4 format (LangChain AWS SDK) - ( - "AWS4-HMAC-SHA256 Credential=Bearer sk-12345678/20260210/us-east-1/bedrock/aws4_request, SignedHeaders=host, Signature=abc123", - "sk-12345678", - "AWS4-HMAC-SHA256 Credential=Bearer sk-12345678/20260210/us-east-1/bedrock/aws4_request, SignedHeaders=host, Signature=abc123", - ), - ], -) -def test_get_api_key_with_custom_litellm_key_header( +def _assert_get_api_key_with_custom_litellm_key_header( custom_litellm_key_header, api_key, passed_in_key ): assert get_api_key( @@ -587,6 +572,49 @@ def test_get_api_key_with_custom_litellm_key_header( ) == (api_key, passed_in_key) +def test_get_api_key_with_custom_litellm_key_header_bearer_prefix(): + token = "sk-" + "1" * 8 + header = f"Bearer {token}" + _assert_get_api_key_with_custom_litellm_key_header( + custom_litellm_key_header=header, api_key=token, passed_in_key=header + ) + + +def test_get_api_key_with_custom_litellm_key_header_basic_prefix(): + token = "sk-" + "1" * 8 + header = f"Basic {token}" + _assert_get_api_key_with_custom_litellm_key_header( + custom_litellm_key_header=header, api_key=token, passed_in_key=header + ) + + +def test_get_api_key_with_custom_litellm_key_header_lowercase_bearer_prefix(): + token = "sk-" + "1" * 8 + header = f"bearer {token}" + _assert_get_api_key_with_custom_litellm_key_header( + custom_litellm_key_header=header, api_key=token, passed_in_key=header + ) + + +def test_get_api_key_with_custom_litellm_key_header_no_prefix(): + token = "sk-" + "1" * 8 + _assert_get_api_key_with_custom_litellm_key_header( + custom_litellm_key_header=token, api_key=token, passed_in_key=token + ) + + +def test_get_api_key_with_custom_litellm_key_header_aws_sigv4(): + """AWS Signature V4 format (LangChain AWS SDK).""" + token = "sk-" + "1" * 8 + header = ( + f"AWS4-HMAC-SHA256 Credential=Bearer {token}/20260210/us-east-1/bedrock/" + "aws4_request, SignedHeaders=host, Signature=abc123" + ) + _assert_get_api_key_with_custom_litellm_key_header( + custom_litellm_key_header=header, api_key=token, passed_in_key=header + ) + + def test_team_metadata_with_tags_flows_through_jwt_auth(): """ Test that team_metadata (specifically tags) flows through JWT authentication.