From 2526b7336d974ab33938a35bff7393c643f88586 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 16 Feb 2026 15:04:34 -0800 Subject: [PATCH] Fixing mapped tests --- .../test_meta_llama_chat_transformation.py | 82 +++++-------------- .../test_publicai_chat_transformation.py | 16 +++- .../test_vertex_ai_rerank_transformation.py | 12 ++- 3 files changed, 41 insertions(+), 69 deletions(-) diff --git a/tests/test_litellm/llms/meta_llama/test_meta_llama_chat_transformation.py b/tests/test_litellm/llms/meta_llama/test_meta_llama_chat_transformation.py index fa605154bb0..7b974aba35c 100644 --- a/tests/test_litellm/llms/meta_llama/test_meta_llama_chat_transformation.py +++ b/tests/test_litellm/llms/meta_llama/test_meta_llama_chat_transformation.py @@ -1,6 +1,5 @@ import os import sys -from unittest.mock import AsyncMock, patch import pytest @@ -47,67 +46,26 @@ def test_map_openai_params(): assert "response_format" in result -@pytest.mark.asyncio -async def test_llama_api_streaming_no_307_error(): - """Test that streaming works without 307 redirect errors due to follow_redirects=True""" +def test_llama_api_streaming_no_307_error(): + """ + Test that the OpenAI-compatible httpx clients use follow_redirects=True. - # Mock the httpx client to simulate a successful streaming response - with patch( - "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" - ) as mock_get_client: - # Create a mock client - mock_client = AsyncMock() - mock_get_client.return_value = mock_client + meta_llama routes through the OpenAI SDK path (BaseOpenAILLM), so the + follow_redirects setting on that SDK's underlying httpx client is what + actually prevents 307 redirect errors for LLaMA API streaming. + """ + from litellm.llms.openai.common_utils import BaseOpenAILLM - # Mock a successful streaming response (not a 307 redirect) - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "text/plain; charset=utf-8"} + # Verify the async httpx client has follow_redirects enabled + async_client = BaseOpenAILLM._get_async_http_client() + assert async_client is not None + assert ( + async_client.follow_redirects is True + ), "Async httpx client should set follow_redirects=True to prevent 307 errors" - # Mock streaming data that would come from a successful request - async def mock_aiter_lines(): - yield 'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"meta_llama/Llama-4-Maverick-17B-128E-Instruct-FP8","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}' - yield 'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"meta_llama/Llama-4-Maverick-17B-128E-Instruct-FP8","choices":[{"index":0,"delta":{"content":" there"},"finish_reason":null}]}' - yield 'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"meta_llama/Llama-4-Maverick-17B-128E-Instruct-FP8","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' - yield "data: [DONE]" - - mock_response.aiter_lines.return_value = mock_aiter_lines() - mock_client.stream.return_value.__aenter__.return_value = mock_response - - # Test the streaming completion - try: - response = await litellm.acompletion( - model="meta_llama/Llama-4-Maverick-17B-128E-Instruct-FP8", - messages=[{"role": "user", "content": "Tell me about yourself"}], - stream=True, - temperature=0.0, - ) - - # Verify we get a CustomStreamWrapper (streaming response) - from litellm.utils import CustomStreamWrapper - - assert isinstance(response, CustomStreamWrapper) - - # Verify the HTTP client was called with follow_redirects=True - mock_client.stream.assert_called_once() - call_kwargs = mock_client.stream.call_args[1] - assert ( - call_kwargs.get("follow_redirects") is True - ), "follow_redirects should be True to prevent 307 errors" - - # Verify the response status is 200 (not 307) - assert ( - mock_response.status_code == 200 - ), "Should get 200 response, not 307 redirect" - - except Exception as e: - # If there's an exception, make sure it's not a 307 error - error_str = str(e) - assert ( - "307" not in error_str - ), f"Should not get 307 redirect error: {error_str}" - - # Still verify that follow_redirects was set correctly - if mock_client.stream.called: - call_kwargs = mock_client.stream.call_args[1] - assert call_kwargs.get("follow_redirects") is True + # Verify the sync httpx client has follow_redirects enabled + sync_client = BaseOpenAILLM._get_sync_http_client() + assert sync_client is not None + assert ( + sync_client.follow_redirects is True + ), "Sync httpx client should set follow_redirects=True to prevent 307 errors" diff --git a/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py b/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py index f6e5e05fe51..cd530cd3b40 100644 --- a/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py +++ b/tests/test_litellm/llms/publicai/test_publicai_chat_transformation.py @@ -7,6 +7,7 @@ PublicAI is an OpenAI-compatible provider with minor customizations. import os import sys +from unittest.mock import patch sys.path.insert( 0, os.path.abspath("../../../../..") @@ -51,9 +52,13 @@ class TestPublicAIConfig: assert result["Authorization"] == f"Bearer {api_key}" assert result["Content-Type"] == "application/json" - def test_get_supported_openai_params(self, config): + @patch("litellm.utils.supports_function_calling", return_value=True) + def test_get_supported_openai_params(self, mock_supports_fc, config): """ - Test that get_supported_openai_params returns correct params + Test that get_supported_openai_params returns correct params. + We mock supports_function_calling because the test model name + 'swiss-ai-apertus' is not in the model registry; this test validates + config behaviour, not registry lookups. """ supported_params = config.get_supported_openai_params(model="swiss-ai-apertus") @@ -66,9 +71,12 @@ class TestPublicAIConfig: # Note: JSON-based configs inherit from OpenAIGPTConfig which includes functions # This is expected behavior for JSON-based providers - def test_map_openai_params_includes_functions(self, config): + @patch("litellm.utils.supports_function_calling", return_value=True) + def test_map_openai_params_includes_functions(self, mock_supports_fc, config): """ - Test that functions parameter is mapped (JSON-based configs don't exclude functions) + Test that functions parameter is mapped (JSON-based configs don't exclude functions). + We mock supports_function_calling because the test model name + 'swiss-ai-apertus' is not in the model registry. """ non_default_params = { "functions": [{"name": "test_function", "description": "Test function"}], diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index fbf5239797f..5e29f927b67 100644 --- a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py @@ -22,6 +22,8 @@ class TestVertexAIRerankTransform: "GOOGLE_APPLICATION_CREDENTIALS", "GOOGLE_CLOUD_PROJECT", "VERTEXAI_PROJECT", + "VERTEXAI_CREDENTIALS", + "VERTEX_AI_CREDENTIALS", "VERTEX_PROJECT", "VERTEX_LOCATION", "VERTEX_AI_PROJECT", @@ -471,16 +473,20 @@ class TestVertexAIRerankTransform: } assert headers == expected_headers - @patch('litellm.llms.vertex_ai.rerank.transformation.VertexAIRerankConfig._ensure_access_token') def test_validate_environment_preserves_optional_params_for_get_complete_url( self, - mock_ensure_access_token, ): """ Validate that calling validate_environment does not remove vertex-specific parameters needed later by get_complete_url. + + Uses instance-level mocking to avoid class-reference issues caused by + importlib.reload(litellm) in conftest.py. """ - mock_ensure_access_token.return_value = ("test-access-token", "project-from-token") + mock_ensure_access_token = MagicMock( + return_value=("test-access-token", "project-from-token") + ) + self.config._ensure_access_token = mock_ensure_access_token optional_params = { "vertex_credentials": "path/to/credentials.json",