diff --git a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index 9fbeb78712e..bfa86647373 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -111,10 +111,13 @@ class TestVertexGemmaCompletion: }, } - # Mock the async HTTP handler + # Mock the async HTTP handler and Vertex authentication with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler" - ) as mock_http_handler: + ) as mock_http_handler, patch( + "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + return_value=("fake-access-token", "PROJECT_ID") + ): mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = mock_vertex_response @@ -197,7 +200,10 @@ class TestVertexGemmaCompletion: with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler" - ) as mock_http_handler: + ) as mock_http_handler, patch( + "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + return_value=("fake-access-token", "test-project") + ): mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = invalid_response @@ -264,7 +270,10 @@ class TestVertexGemmaCompletion: with patch( "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" - ) as mock_get_client: + ) as mock_get_client, patch( + "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + return_value=("fake-access-token", "PROJECT_ID") + ): mock_client = Mock() mock_response = Mock() mock_response.status_code = 200 @@ -354,7 +363,10 @@ class TestVertexGemmaCompletion: with patch( "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" - ) as mock_get_client: + ) as mock_get_client, patch( + "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + return_value=("fake-access-token", "PROJECT_ID") + ): mock_client = Mock() mock_response = Mock() mock_response.status_code = 200