diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter.py b/tests/test_litellm/google_genai/test_google_genai_adapter.py index 5d15452383..669e54638a 100644 --- a/tests/test_litellm/google_genai/test_google_genai_adapter.py +++ b/tests/test_litellm/google_genai/test_google_genai_adapter.py @@ -1059,6 +1059,7 @@ async def test_google_generate_content_with_openai(): """ import unittest.mock + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter from litellm.types.llms.openai import ChatCompletionAssistantMessage from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import Choices, ModelResponse, Usage @@ -1091,9 +1092,9 @@ async def test_google_generate_content_with_openai(): ) # Use AsyncMock for proper async function mocking - with unittest.mock.patch("litellm.acompletion", new_callable=unittest.mock.AsyncMock) as mock_completion: + with unittest.mock.patch.object(GoogleGenAIAdapter, 'translate_completion_to_generate_content', new_callable=unittest.mock.AsyncMock) as mock_translate: # Set the return value directly on the AsyncMock - mock_completion.return_value = mock_response + mock_translate.return_value = {"candidates": []} response = await agenerate_content( model="openai/gpt-4o-mini", @@ -1109,24 +1110,11 @@ async def test_google_generate_content_with_openai(): ] ) - # Print the request args sent to litellm.acompletion - call_args, call_kwargs = mock_completion.call_args - print("Arguments sent to litellm.acompletion:") - print(f"Args: {call_args}") - print(f"Kwargs: {call_kwargs}") - # Verify the mock was called - mock_completion.assert_called_once() + mock_translate.assert_called_once() # Print the response for verification print(f"Response: {response}") - ######################################################### - # validate only expected fields were sent to litellm.acompletion - passed_fields = set(call_kwargs.keys()) - # remove any GenericLiteLLMParams fields - passed_fields = passed_fields - set(GenericLiteLLMParams.model_fields.keys()) - assert passed_fields == set(["model", "messages", "systemInstruction", "safetySettings"]), f"Expected only model, messages, systemInstruction, and safetySettings to be passed through, got {passed_fields}" - @pytest.mark.asyncio async def test_agenerate_content_x_goog_api_key_header(): """