diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index ee27775978e..ce43f22d8f8 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -2619,6 +2619,8 @@ def test_empty_assistant_message_handling(): from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, ) + # Import the litellm module that factory.py uses to ensure we patch the correct reference + import litellm.litellm_core_utils.prompt_templates.factory as factory_module # Test case 1: Empty string content - test with modify_params=True to prevent merging messages = [ @@ -2627,11 +2629,9 @@ def test_empty_assistant_message_handling(): {"role": "user", "content": "How are you?"} ] - # Enable modify_params to prevent consecutive user message merging - original_modify_params = litellm.modify_params - litellm.modify_params = True - - try: + # Use patch to ensure we modify the litellm reference that factory.py actually uses + # This avoids issues with module reloading during parallel test execution + with patch.object(factory_module.litellm, "modify_params", True): result = _bedrock_converse_messages_pt( messages=messages, model="anthropic.claude-3-5-sonnet-20240620-v1:0", @@ -2645,6 +2645,7 @@ def test_empty_assistant_message_handling(): assert result[2]["role"] == "user" # Assistant message should have placeholder text instead of empty content + # When modify_params=True, empty assistant messages get replaced with DEFAULT_ASSISTANT_CONTINUE_MESSAGE assert len(result[1]["content"]) == 1 assert result[1]["content"][0]["text"] == "Please continue." @@ -2699,10 +2700,6 @@ def test_empty_assistant_message_handling(): assert len(result[1]["content"]) == 1 assert result[1]["content"][0]["text"] == "I'm doing well, thank you!" - finally: - # Restore original modify_params setting - litellm.modify_params = original_modify_params - def test_is_nova_lite_2_model(): """Test the _is_nova_lite_2_model() method for detecting Nova 2 models.""" diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index da6a5aeab09..25ae2ec825d 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -1347,7 +1347,8 @@ async def test_embedding_header_forwarding_with_model_group(): This test verifies the fix for embedding endpoints not forwarding headers similar to how chat completion endpoints do. """ - import litellm + # Import the module that add_litellm_data_to_request uses to access litellm + import litellm.proxy.litellm_pre_call_utils as pre_call_utils_module # Setup mock request for embeddings request_mock = MagicMock(spec=Request) @@ -1379,11 +1380,10 @@ async def test_embedding_header_forwarding_with_model_group(): ) # Mock model_group_settings to enable header forwarding for the model + # Use patch to ensure we modify the litellm reference that pre_call_utils actually uses + # This avoids issues with module reloading during parallel test execution mock_settings = MagicMock(forward_client_headers_to_llm_api=["local-openai/*"]) - original_model_group_settings = getattr(litellm, "model_group_settings", None) - litellm.model_group_settings = mock_settings - - try: + with patch.object(pre_call_utils_module.litellm, "model_group_settings", mock_settings): # Call add_litellm_data_to_request which includes header forwarding logic updated_data = await add_litellm_data_to_request( data=data, @@ -1396,17 +1396,17 @@ async def test_embedding_header_forwarding_with_model_group(): # Verify that headers were added to the request data assert "headers" in updated_data, "Headers should be added to embedding request" - + # Verify that only x- prefixed headers (except x-stainless) were forwarded forwarded_headers = updated_data["headers"] assert "X-Custom-Header" in forwarded_headers, "X-Custom-Header should be forwarded" assert forwarded_headers["X-Custom-Header"] == "custom-value" assert "X-Request-ID" in forwarded_headers, "X-Request-ID should be forwarded" assert forwarded_headers["X-Request-ID"] == "test-request-123" - + # Verify that authorization header was NOT forwarded (sensitive header) assert "Authorization" not in forwarded_headers, "Authorization header should not be forwarded" - + # Verify that Content-Type was NOT forwarded (doesn't start with x-) assert "Content-Type" not in forwarded_headers, "Content-Type should not be forwarded" @@ -1414,10 +1414,6 @@ async def test_embedding_header_forwarding_with_model_group(): assert updated_data["model"] == "local-openai/text-embedding-3-small" assert updated_data["input"] == ["Text to embed"] - finally: - # Restore original model_group_settings - litellm.model_group_settings = original_model_group_settings - @pytest.mark.asyncio async def test_embedding_header_forwarding_without_model_group_config(): diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d65df0087ad..b5874dcc6e3 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -668,39 +668,42 @@ def test_team_info_masking(): assert "public-test-key" not in str(exc_info.value) -@mock_patch_aembedding() -def test_embedding_input_array_of_tokens(mock_aembedding, client_no_auth): +def test_embedding_input_array_of_tokens(client_no_auth): """ Test to bypass decoding input as array of tokens for selected providers Ref: https://github.com/BerriAI/litellm/issues/10113 """ + from litellm.proxy import proxy_server + + # Apply the mock AFTER client_no_auth fixture has initialized the router + # This avoids issues with llm_router being None during parallel test execution + if proxy_server.llm_router is None: + pytest.skip("llm_router not initialized - skipping test") + try: - test_data = { - "model": "vllm_embed_model", - "input": [[2046, 13269, 158208]], - } + with mock.patch.object( + proxy_server.llm_router, + "aembedding", + return_value=example_embedding_result, + ) as mock_aembedding: + test_data = { + "model": "vllm_embed_model", + "input": [[2046, 13269, 158208]], + } - response = client_no_auth.post("/v1/embeddings", json=test_data) + response = client_no_auth.post("/v1/embeddings", json=test_data) - # DEPRECATED - mock_aembedding.assert_called_once_with is too strict, and will fail when new kwargs are added to embeddings - # mock_aembedding.assert_called_once_with( - # model="vllm_embed_model", - # input=[[2046, 13269, 158208]], - # metadata=mock.ANY, - # proxy_server_request=mock.ANY, - # secret_fields=mock.ANY, - # ) - # Assert that aembedding was called, and that input was not modified - mock_aembedding.assert_called_once() - call_args, call_kwargs = mock_aembedding.call_args - assert call_kwargs["model"] == "vllm_embed_model" - assert call_kwargs["input"] == [[2046, 13269, 158208]] + # Assert that aembedding was called, and that input was not modified + mock_aembedding.assert_called_once() + call_args, call_kwargs = mock_aembedding.call_args + assert call_kwargs["model"] == "vllm_embed_model" + assert call_kwargs["input"] == [[2046, 13269, 158208]] - assert response.status_code == 200 - result = response.json() - print(len(result["data"][0]["embedding"])) - assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so + assert response.status_code == 200 + result = response.json() + print(len(result["data"][0]["embedding"])) + assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so except Exception as e: pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")