From a7870a902a281e61a4842dbc4bd079fd621a21cf Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 12:15:34 -0700
Subject: [PATCH] test(unified_google_tests): import ReadOnly from
 typing_extensions and cover the Vertex global endpoint

The first commit imported ReadOnly from typing, which only exists on Python 3.13 and up. CircleCI runs this suite on 3.12, so the module failed at import and the job stopped at collection before any of its tests ran. ReadOnly and TypedDict now come from typing_extensions, like the rest of the repo

A new test resolves the Vertex deployment's location from the suite's config with VERTEXAI_LOCATION set to a region, and fails if the vertex_location line is removed

The expected minimum backoff is now derived from litellm's INITIAL_RETRY_DELAY and MAX_RETRY_DELAY, so the test holds when those are overridden through the environment
---
 .../test_google_genai_proxy_test_config.py    | 50 +++++++++++++++----
 1 file changed, 40 insertions(+), 10 deletions(-)

diff --git a/tests/unified_google_tests/test_google_genai_proxy_test_config.py b/tests/unified_google_tests/test_google_genai_proxy_test_config.py
index d84eefb406b..694ec336bac 100644
--- a/tests/unified_google_tests/test_google_genai_proxy_test_config.py
+++ b/tests/unified_google_tests/test_google_genai_proxy_test_config.py
@@ -1,19 +1,26 @@
 import time
 from pathlib import Path
-from typing import Final, ReadOnly, TypedDict
+from typing import Final
 
 import httpx
 import pytest
 import respx
 import yaml
 from pydantic import TypeAdapter
+from typing_extensions import ReadOnly, TypedDict
 
 import litellm
 from litellm import Router
+from litellm.constants import INITIAL_RETRY_DELAY, MAX_RETRY_DELAY
+from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
+from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
 
 CONFIG_PATH: Final = Path(__file__).parent / "google_genai_proxy_test_config.yaml"
+GEMINI_DEPLOYMENT: Final = "gemini-2.5-flash-lite"
+VERTEX_DEPLOYMENT: Final = "vertex-gemini-2.5-flash-lite"
 GEMINI_HOST: Final = "generativelanguage.googleapis.com"
 GEMINI_GENERATE_CONTENT_PATH: Final = "/v1beta/models/gemini-2.5-flash-lite:generateContent"
+VERTEX_GLOBAL_BASE_URL: Final = "https://aiplatform.googleapis.com"
 RESOURCE_EXHAUSTED: Final = {
     "error": {"code": 429, "message": "Resource exhausted. Please try again later.", "status": "RESOURCE_EXHAUSTED"}
 }
@@ -22,7 +29,9 @@ PONG: Final = {
     "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 1, "totalTokenCount": 9},
 }
 CONSECUTIVE_RATE_LIMITS: Final = 3
-MINIMUM_BACKOFF_SECONDS: Final = 0.5 + 1.0 + 2.0
+MINIMUM_BACKOFF_SECONDS: Final = sum(
+    min(INITIAL_RETRY_DELAY * 2**attempt, MAX_RETRY_DELAY) for attempt in range(CONSECUTIVE_RATE_LIMITS)
+)
 
 
 class _Deployment(TypedDict):
@@ -35,14 +44,35 @@ class _ProxyConfig(TypedDict):
     router_settings: ReadOnly[dict[str, dict[str, int]]]
 
 
+def _ci_proxy_config() -> _ProxyConfig:
+    return TypeAdapter(_ProxyConfig).validate_python(yaml.safe_load(CONFIG_PATH.read_text()))
+
+
+def _litellm_params(config: _ProxyConfig, model_name: str) -> dict[str, str]:
+    return next(
+        deployment["litellm_params"] for deployment in config["model_list"] if deployment["model_name"] == model_name
+    )
+
+
 def _router_from_ci_proxy_config() -> Router:
-    config: Final = TypeAdapter(_ProxyConfig).validate_python(yaml.safe_load(CONFIG_PATH.read_text()))
-    gemini_deployments: Final = [
-        {"model_name": deployment["model_name"], "litellm_params": {**deployment["litellm_params"], "api_key": "test"}}
-        for deployment in config["model_list"]
-        if deployment["model_name"] == "gemini-2.5-flash-lite"
-    ]
-    return Router(model_list=gemini_deployments, retry_policy=config["router_settings"]["retry_policy"])
+    config: Final = _ci_proxy_config()
+    return Router(
+        model_list=[
+            {
+                "model_name": GEMINI_DEPLOYMENT,
+                "litellm_params": {**_litellm_params(config, GEMINI_DEPLOYMENT), "api_key": "test"},
+            }
+        ],
+        retry_policy=config["router_settings"]["retry_policy"],
+    )
+
+
+def test_ci_proxy_config_sends_vertex_calls_to_the_global_endpoint(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("VERTEXAI_LOCATION", "us-east5")
+    location: Final = VertexBase.safe_get_vertex_ai_location(_litellm_params(_ci_proxy_config(), VERTEX_DEPLOYMENT))
+
+    assert location == "global"
+    assert get_vertex_base_url(location) == VERTEX_GLOBAL_BASE_URL
 
 
 @pytest.mark.asyncio
@@ -57,7 +87,7 @@ async def test_ci_proxy_config_rides_out_consecutive_429s_with_backoff(
     )
     started: Final = time.monotonic()
     response: Final = await _router_from_ci_proxy_config().agenerate_content(
-        model="gemini-2.5-flash-lite",
+        model=GEMINI_DEPLOYMENT,
         contents=[{"role": "user", "parts": [{"text": "Reply with only the single word: pong"}]}],
     )
     elapsed: Final = time.monotonic() - started