From 3d45ba3edf85b4c400041fa05b9fb64f5f5a1398 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Sun, 15 Mar 2026 18:01:47 -0700
Subject: [PATCH] Fix flaky vertex_ai overhead test by mocking auth and HTTP
 calls

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../test_litellm_overhead.py                  | 74 +++++++++++++++++--
 1 file changed, 68 insertions(+), 6 deletions(-)

diff --git a/tests/litellm_utils_tests/test_litellm_overhead.py b/tests/litellm_utils_tests/test_litellm_overhead.py
index e3472de1848..006fbea8d4b 100644
--- a/tests/litellm_utils_tests/test_litellm_overhead.py
+++ b/tests/litellm_utils_tests/test_litellm_overhead.py
@@ -2,8 +2,10 @@ import json
 import os
 import sys
 import time
+from contextlib import asynccontextmanager, contextmanager
 from datetime import datetime
 from unittest.mock import AsyncMock, patch, MagicMock
+import httpx
 import pytest
 import asyncio
 
@@ -13,6 +15,63 @@ sys.path.insert(
 import litellm
 
 
+# Fake Vertex AI Gemini response for mocking
+FAKE_VERTEX_GEMINI_RESPONSE = {
+    "candidates": [
+        {
+            "content": {
+                "parts": [{"text": "Hello! How can I help you today?"}],
+                "role": "model",
+            },
+            "finishReason": "STOP",
+        }
+    ],
+    "usageMetadata": {
+        "promptTokenCount": 5,
+        "candidatesTokenCount": 8,
+        "totalTokenCount": 13,
+    },
+}
+
+
+def _make_fake_httpx_response(url: str) -> httpx.Response:
+    """Create a fake httpx.Response that looks like a Vertex AI Gemini response."""
+    response = httpx.Response(
+        status_code=200,
+        json=FAKE_VERTEX_GEMINI_RESPONSE,
+        request=httpx.Request("POST", url),
+    )
+    return response
+
+
+@asynccontextmanager
+async def _vertex_ai_mocks():
+    """Context manager that mocks Vertex AI auth and HTTP calls.
+
+    Mocks at the httpx.AsyncClient.send level so that the
+    @track_llm_api_timing decorator on AsyncHTTPHandler.post still runs,
+    preserving the overhead measurement.
+    """
+    fake_response = _make_fake_httpx_response(
+        "https://fake-vertex-endpoint/v1/models/gemini-1.5-flash:generateContent"
+    )
+
+    async def fake_send(self, request, **kwargs):
+        await asyncio.sleep(0.2)  # simulate ~200ms network latency
+        return fake_response
+
+    with patch(
+        "litellm.llms.vertex_ai.vertex_llm_base.VertexBase._ensure_access_token_async",
+        new_callable=AsyncMock,
+        return_value=("Bearer fake-token", "fake-project"),
+    ), patch.object(
+        httpx.AsyncClient,
+        "send",
+        new=fake_send,
+    ):
+        yield
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "model",
@@ -39,16 +98,19 @@ async def test_litellm_overhead_non_streaming(model):
     # Specific cases for models
     #########################################################
     if model == "vertex_ai/gemini-1.5-flash":
-        kwargs["api_base"] = "https://exampleopenaiendpoint-production.up.railway.app/v1/projects/pathrise-convert-1606954137718/locations/us-central1/publishers/google/models/gemini-1.0-pro-vision-001"
-        # warmup call for auth validation on vertex_ai models
-        await litellm.acompletion(**kwargs)
+        kwargs["vertex_project"] = "fake-project"
+        kwargs["vertex_location"] = "us-central1"
     if model == "openai/self_hosted":
         kwargs["api_base"] = "https://exampleopenaiendpoint-production.up.railway.app/"
 
+    async def _run():
+        return await litellm.acompletion(**kwargs)
 
-    response = await litellm.acompletion(
-        **kwargs
-    )
+    if model == "vertex_ai/gemini-1.5-flash":
+        async with _vertex_ai_mocks():
+            response = await _run()
+    else:
+        response = await _run()
     #########################################################
     # End of specific cases for models
     #########################################################