From eec4ed640bf139f46379d70ca85e1f3c03b1f83e Mon Sep 17 00:00:00 2001 From: YutaSaito <36355491+uc4w6c@users.noreply.github.com> Date: Sat, 17 Jan 2026 06:26:18 +0900 Subject: [PATCH] Revert "Stabilise mock tests" --- .../llm_passthrough_endpoints.py | 1 + .../test_responses_background_cost.py | 42 +++-- ...erimental_pass_through_messages_handler.py | 104 ++++++------- .../chat/test_converse_transformation.py | 93 +++++++++++ .../files/test_bedrock_files_integration.py | 146 ++++++++---------- .../huggingface/embedding/test_handler.py | 34 +++- .../files/test_vertex_ai_files_integration.py | 46 ++++++ .../test_openapi_to_mcp_generator.py | 34 ++++ .../guardrails/test_pillar_guardrails.py | 29 +++- .../proxy/test_litellm_pre_call_utils.py | 16 +- tests/test_litellm/proxy/test_proxy_server.py | 39 ++++- tests/test_litellm/test_router.py | 22 +-- 12 files changed, 416 insertions(+), 190 deletions(-) rename tests/test_litellm/{enterprise => integrations}/test_responses_background_cost.py (95%) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 92e37c6408..e48fd22bc8 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -761,6 +761,7 @@ async def handle_bedrock_passthrough_router_model( proxy_logging_obj=proxy_logging_obj, ) + async def handle_bedrock_count_tokens( endpoint: str, request: Request, diff --git a/tests/test_litellm/enterprise/test_responses_background_cost.py b/tests/test_litellm/integrations/test_responses_background_cost.py similarity index 95% rename from tests/test_litellm/enterprise/test_responses_background_cost.py rename to tests/test_litellm/integrations/test_responses_background_cost.py index df694e7adc..6f1e7e9610 100644 --- a/tests/test_litellm/enterprise/test_responses_background_cost.py +++ b/tests/test_litellm/integrations/test_responses_background_cost.py @@ -2,28 +2,14 @@ Integration tests for responses API background cost tracking """ +import asyncio import os -import sys from datetime import datetime from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest -sys.path.insert(0, os.path.abspath("../../..")) - -# Import litellm first to ensure it's in sys.modules before enterprise imports -import litellm # noqa: E402 - -from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse # noqa: E402 - -# Now import enterprise modules -try: - from litellm_enterprise.proxy.common_utils.check_responses_cost import ( # noqa: E402 - CheckResponsesCost, - ) -except ImportError as e: - # Skip all tests in this module if enterprise module is not available - pytest.skip(f"Enterprise module not available: {e}", allow_module_level=True) +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse class TestResponsesBackgroundCostTracking: @@ -298,6 +284,10 @@ class TestCheckResponsesCost: self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router ): """Test CheckResponsesCost initialization""" + from litellm_enterprise.proxy.common_utils.check_responses_cost import ( + CheckResponsesCost, + ) + checker = CheckResponsesCost( proxy_logging_obj=mock_proxy_logging_obj, prisma_client=mock_prisma_client, @@ -313,6 +303,10 @@ class TestCheckResponsesCost: self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router ): """Test polling when there are no jobs""" + from litellm_enterprise.proxy.common_utils.check_responses_cost import ( + CheckResponsesCost, + ) + # Mock find_many to return empty list mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( return_value=[] @@ -340,6 +334,10 @@ class TestCheckResponsesCost: self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router ): """Test polling with a completed job""" + from litellm_enterprise.proxy.common_utils.check_responses_cost import ( + CheckResponsesCost, + ) + # Create a mock job mock_job = MagicMock() mock_job.id = "job-123" @@ -393,6 +391,10 @@ class TestCheckResponsesCost: self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router ): """Test polling with a failed job""" + from litellm_enterprise.proxy.common_utils.check_responses_cost import ( + CheckResponsesCost, + ) + # Create a mock job mock_job = MagicMock() mock_job.id = "job-456" @@ -433,6 +435,10 @@ class TestCheckResponsesCost: self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router ): """Test polling with a job still in progress""" + from litellm_enterprise.proxy.common_utils.check_responses_cost import ( + CheckResponsesCost, + ) + # Create a mock job mock_job = MagicMock() mock_job.id = "job-789" @@ -473,6 +479,10 @@ class TestCheckResponsesCost: self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router ): """Test that errors when querying responses are handled gracefully""" + from litellm_enterprise.proxy.common_utils.check_responses_cost import ( + CheckResponsesCost, + ) + # Create a mock job mock_job = MagicMock() mock_job.id = "job-error" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 5cb2c3cd77..66d62aae1e 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -101,69 +101,55 @@ async def test_bedrock_converse_budget_tokens_preserved(): The bug was that the messages -> completion adapter was converting thinking to reasoning_effort and losing the original budget_tokens value, causing it to use the default (128) instead. """ - import os - client = AsyncHTTPHandler() - # Mock at httpx level for better CI compatibility - with patch("httpx.AsyncClient.post") as mock_httpx_post: - with patch.object(client, "post") as mock_post: - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.headers = {} - mock_response.text = "mock response" - mock_response.json.return_value = { - "output": { - "message": { - "role": "assistant", - "content": [{"text": "4"}] - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 10, - "outputTokens": 5, - "totalTokens": 15 + with patch.object(client, "post") as mock_post: + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.text = "mock response" + mock_response.json.return_value = { + "output": { + "message": { + "role": "assistant", + "content": [{"text": "4"}] } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 10, + "outputTokens": 5, + "totalTokens": 15 } - mock_post.return_value = mock_response - mock_httpx_post.return_value = mock_response - - try: - await messages.acreate( - client=client, - max_tokens=1024, - messages=[{"role": "user", "content": "What is 2+2?"}], - model="bedrock/converse/us.anthropic.claude-sonnet-4-20250514-v1:0", - thinking={ - "budget_tokens": 1024, - "type": "enabled" - }, - ) - except Exception: - pass # Expected due to mock response format - - # Check which mock was called (client.post or httpx.AsyncClient.post) - if mock_post.call_count == 0 and mock_httpx_post.call_count == 0: - # Skip test if neither mock was called (CI environment issue) - if os.getenv("CI") == "true": - pytest.skip("Mock not intercepted in CI environment") - else: - pytest.fail("Expected mock to be called but it wasn't") - - # Use whichever mock was actually called - active_mock = mock_post if mock_post.call_count > 0 else mock_httpx_post - - call_kwargs = active_mock.call_args.kwargs - json_data = call_kwargs.get("json") or json.loads(call_kwargs.get("data", "{}")) - print("Request json: ", json.dumps(json_data, indent=4, default=str)) - - additional_fields = json_data.get("additionalModelRequestFields", {}) - thinking_config = additional_fields.get("thinking", {}) - - assert "thinking" in additional_fields, "thinking parameter should be in additionalModelRequestFields" - assert thinking_config.get("type") == "enabled", "thinking.type should be 'enabled'" - assert thinking_config.get("budget_tokens") == 1024, f"thinking.budget_tokens should be 1024, but got {thinking_config.get('budget_tokens')}" + } + mock_post.return_value = mock_response + + try: + await messages.acreate( + client=client, + max_tokens=1024, + messages=[{"role": "user", "content": "What is 2+2?"}], + model="bedrock/converse/us.anthropic.claude-sonnet-4-20250514-v1:0", + thinking={ + "budget_tokens": 1024, + "type": "enabled" + }, + ) + except Exception: + pass # Expected due to mock response format + + mock_post.assert_called_once() + + call_kwargs = mock_post.call_args.kwargs + json_data = call_kwargs.get("json") or json.loads(call_kwargs.get("data", "{}")) + print("Request json: ", json.dumps(json_data, indent=4, default=str)) + + additional_fields = json_data.get("additionalModelRequestFields", {}) + thinking_config = additional_fields.get("thinking", {}) + + assert "thinking" in additional_fields, "thinking parameter should be in additionalModelRequestFields" + assert thinking_config.get("type") == "enabled", "thinking.type should be 'enabled'" + assert thinking_config.get("budget_tokens") == 1024, f"thinking.budget_tokens should be 1024, but got {thinking_config.get('budget_tokens')}" def test_openai_model_with_thinking_converts_to_reasoning_effort(): 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 763d6964d6..692866f855 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -2610,6 +2610,99 @@ def test_request_metadata_not_provided(): assert "requestMetadata" not in request_data +def test_empty_assistant_message_handling(): + """ + Test that empty assistant messages are handled correctly by replacing + empty or whitespace-only content with a placeholder to prevent AWS Bedrock + Converse API 400 Bad Request errors. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + _bedrock_converse_messages_pt, + ) + + # Test case 1: Empty string content - test with modify_params=True to prevent merging + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": ""}, # Empty content + {"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: + result = _bedrock_converse_messages_pt( + messages=messages, + model="anthropic.claude-3-5-sonnet-20240620-v1:0", + llm_provider="bedrock_converse" + ) + + # Should have 3 messages: user, assistant (with placeholder), user + assert len(result) == 3 + assert result[0]["role"] == "user" + assert result[1]["role"] == "assistant" + assert result[2]["role"] == "user" + + # Assistant message should have placeholder text instead of empty content + assert len(result[1]["content"]) == 1 + assert result[1]["content"][0]["text"] == "Please continue." + + # Test case 2: Whitespace-only content + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": " "}, # Whitespace-only content + {"role": "user", "content": "How are you?"} + ] + + result = _bedrock_converse_messages_pt( + messages=messages, + model="anthropic.claude-3-5-sonnet-20240620-v1:0", + llm_provider="bedrock_converse" + ) + + # Assistant message should have placeholder text instead of whitespace + assert len(result[1]["content"]) == 1 + assert result[1]["content"][0]["text"] == "Please continue." + + # Test case 3: Empty list content + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": [{"type": "text", "text": ""}]}, # Empty text in list + {"role": "user", "content": "How are you?"} + ] + + result = _bedrock_converse_messages_pt( + messages=messages, + model="anthropic.claude-3-5-sonnet-20240620-v1:0", + llm_provider="bedrock_converse" + ) + + # Assistant message should have placeholder text instead of empty text + assert len(result[1]["content"]) == 1 + assert result[1]["content"][0]["text"] == "Please continue." + + # Test case 4: Normal content should not be affected + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "I'm doing well, thank you!"}, # Normal content + {"role": "user", "content": "How are you?"} + ] + + result = _bedrock_converse_messages_pt( + messages=messages, + model="anthropic.claude-3-5-sonnet-20240620-v1:0", + llm_provider="bedrock_converse" + ) + + # Assistant message should keep original content + 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/llms/bedrock/files/test_bedrock_files_integration.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py index 983ad73980..37a0daa1d5 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py @@ -21,51 +21,43 @@ class TestBedrockFilesIntegration: file_id = "s3://test-bucket/test-file.jsonl" expected_content = b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}' - # Mock AWS credentials - with patch.dict( - "os.environ", - { - "AWS_ACCESS_KEY_ID": "test-access-key", - "AWS_SECRET_ACCESS_KEY": "test-secret-key", - }, - ): - # Mock the bedrock_files_instance.file_content method - with patch( - "litellm.files.main.bedrock_files_instance.file_content", - new_callable=AsyncMock, - ) as mock_file_content: - # Create a mock HttpxBinaryResponseContent response - import httpx + # Mock the bedrock_files_instance.file_content method + with patch( + "litellm.files.main.bedrock_files_instance.file_content", + new_callable=AsyncMock, + ) as mock_file_content: + # Create a mock HttpxBinaryResponseContent response + import httpx - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request( - method="GET", url="s3://test-bucket/test-file.jsonl" - ), - ) - mock_file_content.return_value = HttpxBinaryResponseContent( - response=mock_response - ) + mock_response = httpx.Response( + status_code=200, + content=expected_content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request( + method="GET", url="s3://test-bucket/test-file.jsonl" + ), + ) + mock_file_content.return_value = HttpxBinaryResponseContent( + response=mock_response + ) - # Call litellm.afile_content - result = await litellm.afile_content( - file_id=file_id, - custom_llm_provider="bedrock", - aws_region_name="us-west-2", - ) + # Call litellm.afile_content + result = await litellm.afile_content( + file_id=file_id, + custom_llm_provider="bedrock", + aws_region_name="us-west-2", + ) - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - assert result.response.status_code == 200 + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + assert result.response.status_code == 200 - # Verify the mock was called with correct parameters - mock_file_content.assert_called_once() - call_kwargs = mock_file_content.call_args.kwargs - assert call_kwargs["_is_async"] is True - assert call_kwargs["file_content_request"]["file_id"] == file_id + # Verify the mock was called with correct parameters + mock_file_content.assert_called_once() + call_kwargs = mock_file_content.call_args.kwargs + assert call_kwargs["_is_async"] is True + assert call_kwargs["file_content_request"]["file_id"] == file_id @pytest.mark.asyncio async def test_litellm_afile_content_bedrock_provider_with_unified_file_id(self): @@ -80,47 +72,39 @@ class TestBedrockFilesIntegration: expected_content = b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}' - # Mock AWS credentials - with patch.dict( - "os.environ", - { - "AWS_ACCESS_KEY_ID": "test-access-key", - "AWS_SECRET_ACCESS_KEY": "test-secret-key", - }, - ): - # Mock the bedrock_files_instance.file_content method - with patch( - "litellm.files.main.bedrock_files_instance.file_content", - new_callable=AsyncMock, - ) as mock_file_content: - # Create a mock HttpxBinaryResponseContent response - import httpx + # Mock the bedrock_files_instance.file_content method + with patch( + "litellm.files.main.bedrock_files_instance.file_content", + new_callable=AsyncMock, + ) as mock_file_content: + # Create a mock HttpxBinaryResponseContent response + import httpx - mock_response = httpx.Response( - status_code=200, - content=expected_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url=s3_uri), - ) - mock_file_content.return_value = HttpxBinaryResponseContent( - response=mock_response - ) + mock_response = httpx.Response( + status_code=200, + content=expected_content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request(method="GET", url=s3_uri), + ) + mock_file_content.return_value = HttpxBinaryResponseContent( + response=mock_response + ) - # Call litellm.afile_content with unified file ID - result = await litellm.afile_content( - file_id=encoded_file_id, - custom_llm_provider="bedrock", - aws_region_name="us-west-2", - ) + # Call litellm.afile_content with unified file ID + result = await litellm.afile_content( + file_id=encoded_file_id, + custom_llm_provider="bedrock", + aws_region_name="us-west-2", + ) - # Verify the result - assert isinstance(result, HttpxBinaryResponseContent) - assert result.response.content == expected_content - assert result.response.status_code == 200 + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + assert result.response.status_code == 200 - # Verify the mock was called - the handler should extract S3 URI from unified file ID - mock_file_content.assert_called_once() - call_kwargs = mock_file_content.call_args.kwargs - assert call_kwargs["_is_async"] is True - # The handler extracts S3 URI from the unified file ID - assert call_kwargs["file_content_request"]["file_id"] == encoded_file_id + # Verify the mock was called - the handler should extract S3 URI from unified file ID + mock_file_content.assert_called_once() + call_kwargs = mock_file_content.call_args.kwargs + assert call_kwargs["_is_async"] is True + # The handler extracts S3 URI from the unified file ID + assert call_kwargs["file_content_request"]["file_id"] == encoded_file_id diff --git a/tests/test_litellm/llms/huggingface/embedding/test_handler.py b/tests/test_litellm/llms/huggingface/embedding/test_handler.py index b768bee403..f6bc983df0 100644 --- a/tests/test_litellm/llms/huggingface/embedding/test_handler.py +++ b/tests/test_litellm/llms/huggingface/embedding/test_handler.py @@ -41,12 +41,8 @@ def mock_embedding_async_http_handler(): class TestHuggingFaceEmbedding: @pytest.fixture(autouse=True) def setup(self, mock_embedding_http_handler, mock_embedding_async_http_handler): - # Mock both sync and async versions of get_hf_task functions self.mock_get_task_patcher = patch("litellm.llms.huggingface.embedding.handler.get_hf_task_embedding_for_model") - self.mock_get_task_async_patcher = patch("litellm.llms.huggingface.embedding.handler.async_get_hf_task_embedding_for_model", new_callable=AsyncMock) - self.mock_get_task = self.mock_get_task_patcher.start() - self.mock_get_task_async = self.mock_get_task_async_patcher.start() def mock_get_task_side_effect(model, task_type, api_base): if task_type is not None: @@ -54,7 +50,6 @@ class TestHuggingFaceEmbedding: return "sentence-similarity" self.mock_get_task.side_effect = mock_get_task_side_effect - self.mock_get_task_async.side_effect = mock_get_task_side_effect self.model = "huggingface/BAAI/bge-m3" self.mock_http = mock_embedding_http_handler @@ -64,7 +59,6 @@ class TestHuggingFaceEmbedding: yield self.mock_get_task_patcher.stop() - self.mock_get_task_async_patcher.stop() def test_input_type_preserved_in_optional_params(self): input_text = ["hello world"] @@ -87,3 +81,31 @@ class TestHuggingFaceEmbedding: # Should NOT have sentence-similarity format assert "source_sentence" not in str(request_data) assert "sentences" not in str(request_data) + + def test_embedding_with_sentence_similarity_task(self): + """Test embedding when task type is sentence-similarity (requires 2+ sentences)""" + + similarity_response = { + "similarities": [[0, 0.9], [1, 0.8]] + } + + self.mock_http.return_value.json.return_value = similarity_response + + # Test with 2+ sentences (required for sentence-similarity) + input_text = ["This is the source sentence", "This is sentence one", "This is sentence two"] + + response = litellm.embedding( + model=self.model, + input=input_text, + # Use the model's natural task type (sentence-similarity) + ) + + self.mock_http.assert_called_once() + post_call_args = self.mock_http.call_args + request_data = json.loads(post_call_args[1]["data"]) + + assert "inputs" in request_data + assert "source_sentence" in request_data["inputs"] + assert "sentences" in request_data["inputs"] + assert request_data["inputs"]["source_sentence"] == input_text[0] + assert request_data["inputs"]["sentences"] == input_text[1:] \ No newline at end of file diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py index 50ad3920cb..723594dc39 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py @@ -12,7 +12,53 @@ from litellm.types.llms.openai import HttpxBinaryResponseContent class TestVertexAIFilesIntegration: """Test integration of Vertex AI files with main litellm API""" + @pytest.mark.asyncio + async def test_litellm_afile_content_vertex_ai_provider(self): + """Test litellm.afile_content with vertex_ai provider""" + file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" + expected_content = b"test file content" + # Mock the vertex_ai_files_instance.file_content method + with patch( + "litellm.files.main.vertex_ai_files_instance.file_content", + new_callable=AsyncMock, + ) as mock_file_content: + # Create a mock HttpxBinaryResponseContent response + import httpx + + mock_response = httpx.Response( + status_code=200, + content=expected_content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request( + method="GET", url="gs://test-bucket/test-file.txt" + ), + ) + mock_file_content.return_value = HttpxBinaryResponseContent( + response=mock_response + ) + + # Call litellm.afile_content + result = await litellm.afile_content( + file_id=file_id, + custom_llm_provider="vertex_ai", + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, + ) + + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + assert result.response.status_code == 200 + + # Verify the mock was called with correct parameters + mock_file_content.assert_called_once() + call_kwargs = mock_file_content.call_args.kwargs + assert call_kwargs["_is_async"] is True + assert call_kwargs["file_content_request"]["file_id"] == file_id + assert call_kwargs["vertex_project"] == "test-project" + assert call_kwargs["vertex_location"] == "us-central1" def test_litellm_file_content_vertex_ai_provider(self): """Test litellm.file_content with vertex_ai provider (sync)""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 488f26cdca..573e095606 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -75,6 +75,40 @@ class TestCreateToolFunction: call_args[0][0] ) + @pytest.mark.asyncio + async def test_leading_digit_parameter(self): + """Test function with parameter starting with digit (e.g., 2fa-code).""" + operation = { + "parameters": [ + { + "name": "2fa-code", + "in": "query", + "required": False, + "schema": {"type": "string"}, + } + ] + } + + func = create_tool_function( + path="/verify", + method="post", + operation=operation, + base_url="https://api.example.com", + ) + + assert callable(func) + + with patch(GET_ASYNC_CLIENT_TARGET) as mock_client: + async_client = _create_mock_client("post", "verified") + mock_client.return_value = async_client + + result = await func(**{"2fa-code": "123456"}) + assert result == "verified" + + # Verify query parameter was included + call_args = async_client.post.call_args + assert call_args[1]["params"]["2fa-code"] == "123456" + @pytest.mark.asyncio async def test_dot_in_parameter_name(self): """Test function with dot in parameter name (e.g., user.name).""" diff --git a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py index 681caf9716..0607b0de98 100644 --- a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py @@ -8,7 +8,7 @@ and following LiteLLM testing patterns and best practices. # Standard library imports import os import sys -from typing import Dict, Any +from typing import Dict from unittest.mock import Mock, patch # Add parent directory to path for imports @@ -43,6 +43,33 @@ from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 # ============================================================================ +@pytest.fixture(scope="function", autouse=True) +def setup_and_teardown(): + """ + Standard LiteLLM fixture that reloads litellm before every function + to speed up testing by removing callbacks being chained. + """ + import importlib + import asyncio + + # Reload litellm to ensure clean state + importlib.reload(litellm) + + # Set up async loop + loop = asyncio.get_event_loop_policy().new_event_loop() + asyncio.set_event_loop(loop) + + # Set up litellm state + litellm.set_verbose = True + litellm.guardrail_name_config_map = {} + + yield + + # Teardown + loop.close() + asyncio.set_event_loop(None) + + @pytest.fixture def env_setup(monkeypatch): """Fixture to set up environment variables for testing.""" 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 cc7ffeb0b6..133fc07d34 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -1393,23 +1393,21 @@ async def test_embedding_header_forwarding_with_model_group(): version="test-version", ) - # Verify that headers were added to the request metadata - assert "metadata" in updated_data, "Metadata should be added to embedding request" - assert "headers" in updated_data["metadata"], "Headers should be added to embedding request metadata" + # 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["metadata"]["headers"] + 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 is present in metadata (not filtered out at this level) - # Note: The metadata headers contain all original headers for logging/tracking purposes - assert "Authorization" in forwarded_headers, "Authorization header should be in metadata headers" + # 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 is present (it's included in metadata headers) - assert "Content-Type" in forwarded_headers, "Content-Type should be in metadata headers" + # 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" # Verify original data fields are preserved assert updated_data["model"] == "local-openai/text-embedding-3-small" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d14ac5cf33..751a903387 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -55,7 +55,7 @@ example_embedding_result = { def mock_patch_aembedding(): return mock.patch( - "litellm.aembedding", + "litellm.proxy.proxy_server.llm_router.aembedding", return_value=example_embedding_result, ) @@ -668,6 +668,43 @@ 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): + """ + Test to bypass decoding input as array of tokens for selected providers + + Ref: https://github.com/BerriAI/litellm/issues/10113 + """ + try: + test_data = { + "model": "vllm_embed_model", + "input": [[2046, 13269, 158208]], + } + + 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 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)}") + + @pytest.mark.asyncio async def test_get_all_team_models(): """ diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 12fc65d8b0..7201b96158 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1231,30 +1231,18 @@ async def test_acompletion_streaming_disable_fallbacks_midstream(): return self async def __anext__(self): - if self.index == self.error_after_index: - raise self.error if self.index >= len(self.items): raise StopAsyncIteration + if self.index == self.error_after_index: + raise self.error item = self.items[self.index] self.index += 1 self.chunks.append(item) return item - # Create properly structured mock chunks using ModelResponse - from litellm.types.utils import Delta, ModelResponse, StreamingChoices - - mock_chunk = ModelResponse( - id="chatcmpl-123", - choices=[ - StreamingChoices( - index=0, delta=Delta(content="Hello", role="assistant"), finish_reason=None - ) - ], - created=1234567890, - model="gpt-4", - object="chat.completion.chunk", - ) - mock_chunks = [mock_chunk] + mock_chunks = [ + MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]), + ] mock_error_response = AsyncIteratorWithError( mock_chunks, 1, error_with_original