mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-29 14:21:40 +00:00
0c006794f1
* litellm_fix_mapped_tests_core: fix test isolation and mock injection issues ## Problem Four tests in litellm_mapped_tests_core were failing: 1. test_register_model_with_scientific_notation - KeyError due to test isolation issues 2. test_search_uses_registry_credentials - Mock not being called due to incorrect patch path 3. test_send_email_missing_api_key - Real API calls despite mocking 4. test_stream_transformation_error_sync - Mock not effective, real API called ## Solution ### test_register_model_with_scientific_notation - Use unique model name to avoid conflicts with other tests - Clear LRU caches before test to prevent stale data - Clean up model_cost entry after test ### test_search_uses_registry_credentials - Use patch.object() on the actual base_llm_http_handler instance - String-based patching for instance methods can fail; direct object patching is more reliable ### test_send_email_missing_api_key - Directly inject mock HTTP client into logger instance - This bypasses any caching issues that could cause the fixture mock to be ineffective ### test_stream_transformation_error_sync - Patch litellm.completion directly instead of the handler module's litellm reference - This ensures the mock is effective regardless of import order ## Regression These tests were affected by LRU caching added in #19606 and HTTP client caching. * fix(test): use patch.object for container API tests to fix mock injection ## Problem test_retrieve_container_basic tests were failing because mocks weren't being applied correctly. The tests used string-based patching: patch('litellm.containers.main.base_llm_http_handler') But base_llm_http_handler is imported at module level, so the mock wasn't intercepting the actual handler calls, resulting in real HTTP requests to OpenAI API. ## Solution Use patch.object() to directly mock methods on the imported handler instance. Import base_llm_http_handler in the test file and patch like: patch.object(base_llm_http_handler, 'container_retrieve_handler', ...) This ensures the mock is applied to the actual object being used, regardless of import order or caching. * fix(test): add missing Prometheus metric labels to test_proxy_failure_metrics Add client_ip, user_agent, model_id labels to expected metric patterns. These labels were added in PRs #19717 and #19678 but test wasn't updated. * fix(test_resend_email): use direct mock injection for all email tests Extend the mock injection pattern used in test_send_email_missing_api_key to all other tests in the file: - test_send_email_success - test_send_email_multiple_recipients Instead of relying on fixture-based patching and respx mocks which can fail due to import order and caching issues, directly inject the mock HTTP client into the logger instance. This ensures mocks are always used regardless of test execution order. * fix(test): use patch.object for image_edit and vector_store tests - test_image_edit_merges_headers_and_extra_headers: import base_llm_http_handler and use patch.object instead of string path patching - test_search_uses_registry_credentials: import module and patch via module.base_llm_http_handler to ensure we patch the right instance --------- Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com>
335 lines
13 KiB
Python
335 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test to verify the Google GenAI generate_content handler functionality
|
|
"""
|
|
import json
|
|
import os
|
|
import sys
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../../../..")
|
|
) # Adds the parent directory to the system path
|
|
|
|
import litellm
|
|
from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler
|
|
from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
|
|
from litellm.types.utils import ModelResponse
|
|
|
|
|
|
def test_non_stream_response_when_stream_requested_sync():
|
|
"""
|
|
Test that when a non-stream response is returned but streaming was requested,
|
|
the sync handler correctly transforms it to generate_content format.
|
|
"""
|
|
from litellm.types.utils import Choices
|
|
|
|
# Mock a non-stream response (ModelResponse with valid choices)
|
|
mock_response = ModelResponse(
|
|
id="test-123",
|
|
choices=[
|
|
Choices(
|
|
index=0,
|
|
message={
|
|
"role": "assistant",
|
|
"content": "Hello, world!"
|
|
},
|
|
finish_reason="stop"
|
|
)
|
|
],
|
|
created=1234567890,
|
|
model="gpt-3.5-turbo",
|
|
object="chat.completion"
|
|
)
|
|
|
|
# Create an instance of the adapter
|
|
adapter = GoogleGenAIAdapter()
|
|
|
|
# Test the adapter's translate_completion_to_generate_content method directly
|
|
result = adapter.translate_completion_to_generate_content(mock_response)
|
|
|
|
# Verify the result is a valid Google GenAI format response
|
|
assert "candidates" in result
|
|
assert isinstance(result["candidates"], list)
|
|
assert len(result["candidates"]) > 0
|
|
candidate = result["candidates"][0]
|
|
assert "content" in candidate
|
|
assert "parts" in candidate["content"]
|
|
assert isinstance(candidate["content"]["parts"], list)
|
|
assert len(candidate["content"]["parts"]) > 0
|
|
assert "text" in candidate["content"]["parts"][0]
|
|
assert candidate["content"]["parts"][0]["text"] == "Hello, world!"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_stream_response_when_stream_requested_async():
|
|
"""
|
|
Test that when a non-stream response is returned but streaming was requested,
|
|
the async handler correctly transforms it to generate_content format.
|
|
"""
|
|
from litellm.types.utils import Choices
|
|
|
|
# Mock a non-stream response (ModelResponse with valid choices)
|
|
mock_response = ModelResponse(
|
|
id="test-123",
|
|
choices=[
|
|
Choices(
|
|
index=0,
|
|
message={
|
|
"role": "assistant",
|
|
"content": "Hello, world!"
|
|
},
|
|
finish_reason="stop"
|
|
)
|
|
],
|
|
created=1234567890,
|
|
model="gpt-3.5-turbo",
|
|
object="chat.completion"
|
|
)
|
|
|
|
# Create an instance of the adapter
|
|
adapter = GoogleGenAIAdapter()
|
|
|
|
# Test the adapter's translate_completion_to_generate_content method directly
|
|
result = adapter.translate_completion_to_generate_content(mock_response)
|
|
|
|
# Verify the result is a valid Google GenAI format response
|
|
assert "candidates" in result
|
|
assert isinstance(result["candidates"], list)
|
|
assert len(result["candidates"]) > 0
|
|
candidate = result["candidates"][0]
|
|
assert "content" in candidate
|
|
assert "parts" in candidate["content"]
|
|
assert isinstance(candidate["content"]["parts"], list)
|
|
assert len(candidate["content"]["parts"]) > 0
|
|
assert "text" in candidate["content"]["parts"][0]
|
|
assert candidate["content"]["parts"][0]["text"] == "Hello, world!"
|
|
|
|
|
|
def test_stream_response_when_stream_requested_sync():
|
|
"""
|
|
Test that when a stream response is returned and streaming was requested,
|
|
the sync handler correctly transforms it to generate_content streaming format.
|
|
"""
|
|
# Mock a stream response
|
|
mock_stream = MagicMock()
|
|
mock_stream.__iter__ = MagicMock(return_value=iter([]))
|
|
|
|
# Mock the GoogleGenAIAdapter's translate_completion_output_params_streaming method
|
|
with patch.object(
|
|
GoogleGenAIAdapter,
|
|
"translate_completion_output_params_streaming",
|
|
return_value=mock_stream
|
|
) as mock_translate:
|
|
with patch("litellm.completion", return_value=mock_stream):
|
|
# Call the handler with stream=True
|
|
result = GenerateContentToCompletionHandler.generate_content_handler(
|
|
model="gemini-pro",
|
|
contents=[{"role": "user", "parts": [{"text": "Hello"}]}],
|
|
litellm_params={}, # Empty dict for params
|
|
stream=True
|
|
)
|
|
|
|
# Verify that translate_completion_output_params_streaming was called
|
|
mock_translate.assert_called_once_with(mock_stream)
|
|
# Verify the result is the transformed stream
|
|
assert result == mock_stream
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_response_when_stream_requested_async():
|
|
"""
|
|
Test that when a stream response is returned and streaming was requested,
|
|
the async handler correctly transforms it to generate_content streaming format.
|
|
"""
|
|
# Mock a stream response
|
|
mock_stream = MagicMock()
|
|
mock_stream.__aiter__ = AsyncMock(return_value=iter([])) # Return an empty async iterator
|
|
|
|
# Mock the GoogleGenAIAdapter's translate_completion_output_params_streaming method
|
|
with patch.object(
|
|
GoogleGenAIAdapter,
|
|
"translate_completion_output_params_streaming",
|
|
return_value=mock_stream
|
|
) as mock_translate:
|
|
with patch("litellm.acompletion", return_value=mock_stream):
|
|
# Call the handler with stream=True
|
|
result = await GenerateContentToCompletionHandler.async_generate_content_handler(
|
|
model="gemini-pro",
|
|
contents=[{"role": "user", "parts": [{"text": "Hello"}]}],
|
|
litellm_params={}, # Empty dict for params
|
|
stream=True
|
|
)
|
|
|
|
# Verify that translate_completion_output_params_streaming was called
|
|
mock_translate.assert_called_once_with(mock_stream)
|
|
# Verify the result is the transformed stream
|
|
assert result == mock_stream
|
|
|
|
|
|
def test_stream_transformation_error_sync():
|
|
"""
|
|
Test that when a stream transformation fails, the sync handler raises a ValueError.
|
|
"""
|
|
# Mock a stream response
|
|
mock_stream = MagicMock()
|
|
mock_stream.__iter__ = MagicMock(return_value=iter([]))
|
|
|
|
# Mock the GoogleGenAIAdapter's translate_completion_output_params_streaming method to return None
|
|
with patch.object(
|
|
GoogleGenAIAdapter,
|
|
"translate_completion_output_params_streaming",
|
|
return_value=None
|
|
):
|
|
# Patch litellm.completion directly to prevent real API calls
|
|
with patch("litellm.completion", return_value=mock_stream):
|
|
# Call the handler with stream=True and expect a ValueError
|
|
with pytest.raises(ValueError, match="Failed to transform streaming response"):
|
|
GenerateContentToCompletionHandler.generate_content_handler(
|
|
model="gemini-pro",
|
|
contents=[{"role": "user", "parts": [{"text": "Hello"}]}],
|
|
litellm_params={}, # Empty dict for params
|
|
stream=True
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_transformation_error_async():
|
|
"""
|
|
Test that when a stream transformation fails, the async handler raises a ValueError.
|
|
"""
|
|
# Mock a stream response
|
|
mock_stream = MagicMock()
|
|
mock_stream.__aiter__ = AsyncMock(return_value=mock_stream)
|
|
|
|
# Mock the GoogleGenAIAdapter's translate_completion_output_params_streaming method to return None
|
|
with patch.object(
|
|
GoogleGenAIAdapter,
|
|
"translate_completion_output_params_streaming",
|
|
return_value=None
|
|
):
|
|
# Mock litellm.acompletion at the module level where it's imported
|
|
# We need to patch it in the handler module, not in litellm itself
|
|
with patch("litellm.google_genai.adapters.handler.litellm") as mock_litellm:
|
|
# Use AsyncMock for async function
|
|
mock_litellm.acompletion = AsyncMock(return_value=mock_stream)
|
|
# Call the handler with stream=True and expect a ValueError
|
|
with pytest.raises(ValueError, match="Failed to transform streaming response"):
|
|
await GenerateContentToCompletionHandler.async_generate_content_handler(
|
|
model="gemini-pro",
|
|
contents=[{"role": "user", "parts": [{"text": "Hello"}]}],
|
|
litellm_params={}, # Empty dict for params
|
|
stream=True
|
|
)
|
|
|
|
|
|
def test_citation_metadata_transformation():
|
|
"""
|
|
Test that citationMetadata.citationSources is properly transformed to citationMetadata.citations
|
|
to avoid Pydantic validation errors.
|
|
"""
|
|
from unittest.mock import MagicMock
|
|
|
|
import httpx
|
|
|
|
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
|
from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig
|
|
|
|
# Create a mock response with citationMetadata.citationSources (the problematic format)
|
|
mock_response_data = {
|
|
"candidates": [
|
|
{
|
|
"content": {
|
|
"parts": [
|
|
{
|
|
"text": "This is a video analysis response with citation metadata."
|
|
}
|
|
],
|
|
"role": "model"
|
|
},
|
|
"finishReason": "STOP",
|
|
"index": 0,
|
|
"safetyRatings": [],
|
|
"citationMetadata": {
|
|
"citationSources": [
|
|
{
|
|
"startIndex": 5848,
|
|
"endIndex": 5900,
|
|
"uri": "https://example.com/video-source",
|
|
"license": "MIT",
|
|
"title": "Video Analysis Source",
|
|
"publicationDate": "2024-01-15"
|
|
},
|
|
{
|
|
"startIndex": 6200,
|
|
"endIndex": 6250,
|
|
"uri": "https://another-source.com/reference",
|
|
"license": "CC-BY",
|
|
"title": "Another Reference",
|
|
"publicationDate": "2024-02-01"
|
|
}
|
|
]
|
|
}
|
|
}
|
|
],
|
|
"usageMetadata": {
|
|
"promptTokenCount": 150,
|
|
"candidatesTokenCount": 200,
|
|
"totalTokenCount": 350
|
|
},
|
|
"responseId": "test-response-123"
|
|
}
|
|
|
|
# Create mock httpx response
|
|
mock_httpx_response = MagicMock(spec=httpx.Response)
|
|
mock_httpx_response.json.return_value = mock_response_data
|
|
mock_httpx_response.status_code = 200
|
|
mock_httpx_response.headers = {}
|
|
|
|
# Create logging object
|
|
logging_obj = LiteLLMLoggingObj(
|
|
model="gemini-2.5-flash",
|
|
messages=[],
|
|
stream=False,
|
|
call_type="generate_content",
|
|
start_time=1234567890,
|
|
litellm_call_id="test-call-123",
|
|
function_id="test-function-123"
|
|
)
|
|
|
|
# Create GoogleGenAI config
|
|
config = GoogleGenAIConfig()
|
|
|
|
# Test the transformation
|
|
try:
|
|
result = config.transform_generate_content_response(
|
|
model="gemini-2.5-flash",
|
|
raw_response=mock_httpx_response,
|
|
logging_obj=logging_obj
|
|
)
|
|
|
|
# Verify the transformation worked
|
|
assert result is not None
|
|
|
|
# Check that citationSources was transformed to citations
|
|
if hasattr(result, 'candidates') and result.candidates:
|
|
candidate = result.candidates[0]
|
|
if hasattr(candidate, 'citationMetadata') and candidate.citationMetadata:
|
|
# The citationMetadata should now have 'citations' instead of 'citationSources'
|
|
citation_metadata = candidate.citationMetadata
|
|
|
|
# Check that citations field exists
|
|
assert hasattr(citation_metadata, 'citations'), "citations field should exist after transformation"
|
|
|
|
# Verify the citations data is preserved
|
|
if hasattr(citation_metadata, 'citations') and citation_metadata.citations:
|
|
assert len(citation_metadata.citations) == 2, "Should have 2 citations"
|
|
assert citation_metadata.citations[0]['uri'] == "https://example.com/video-source"
|
|
assert citation_metadata.citations[1]['uri'] == "https://another-source.com/reference"
|
|
|
|
print("✅ Citation metadata transformation test passed!")
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Citation metadata transformation failed: {e}") |