From 47d06d9fdd5973ea2e72daec07b1a38bae24b2bd Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 11:56:02 -0700
Subject: [PATCH] test(unified_google_tests): use the Vertex global endpoint
 and retry 429s with backoff

The google_generate_content_endpoint_testing job went red on main when us-central1 ran out of shared gemini-2.5-flash-lite capacity for a few hours. The suite's proxy config now sends the Vertex deployment to the global endpoint and retries rate limit errors 5 times with exponential backoff, and a regression test pins that the config rides out 3 consecutive 429s
---
 .../google_genai_proxy_test_config.yaml       |  5 ++
 .../test_google_genai_proxy_test_config.py    | 67 +++++++++++++++++++
 2 files changed, 72 insertions(+)
 create mode 100644 tests/unified_google_tests/test_google_genai_proxy_test_config.py

diff --git a/tests/unified_google_tests/google_genai_proxy_test_config.yaml b/tests/unified_google_tests/google_genai_proxy_test_config.yaml
index 9913c05d434..64a83ef3d81 100644
--- a/tests/unified_google_tests/google_genai_proxy_test_config.yaml
+++ b/tests/unified_google_tests/google_genai_proxy_test_config.yaml
@@ -7,6 +7,11 @@ model_list:
   - model_name: vertex-gemini-2.5-flash-lite
     litellm_params:
       model: vertex_ai/gemini-2.5-flash-lite
+      vertex_location: global
+
+router_settings:
+  retry_policy:
+    RateLimitErrorRetries: 5
 
 general_settings:
   master_key: sk-1234
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
new file mode 100644
index 00000000000..d84eefb406b
--- /dev/null
+++ b/tests/unified_google_tests/test_google_genai_proxy_test_config.py
@@ -0,0 +1,67 @@
+import time
+from pathlib import Path
+from typing import Final, ReadOnly, TypedDict
+
+import httpx
+import pytest
+import respx
+import yaml
+from pydantic import TypeAdapter
+
+import litellm
+from litellm import Router
+
+CONFIG_PATH: Final = Path(__file__).parent / "google_genai_proxy_test_config.yaml"
+GEMINI_HOST: Final = "generativelanguage.googleapis.com"
+GEMINI_GENERATE_CONTENT_PATH: Final = "/v1beta/models/gemini-2.5-flash-lite:generateContent"
+RESOURCE_EXHAUSTED: Final = {
+    "error": {"code": 429, "message": "Resource exhausted. Please try again later.", "status": "RESOURCE_EXHAUSTED"}
+}
+PONG: Final = {
+    "candidates": [{"content": {"role": "model", "parts": [{"text": "pong"}]}, "finishReason": "STOP"}],
+    "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 1, "totalTokenCount": 9},
+}
+CONSECUTIVE_RATE_LIMITS: Final = 3
+MINIMUM_BACKOFF_SECONDS: Final = 0.5 + 1.0 + 2.0
+
+
+class _Deployment(TypedDict):
+    model_name: ReadOnly[str]
+    litellm_params: ReadOnly[dict[str, str]]
+
+
+class _ProxyConfig(TypedDict):
+    model_list: ReadOnly[list[_Deployment]]
+    router_settings: ReadOnly[dict[str, dict[str, int]]]
+
+
+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"])
+
+
+@pytest.mark.asyncio
+async def test_ci_proxy_config_rides_out_consecutive_429s_with_backoff(
+    respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    litellm.in_memory_llm_clients_cache.flush_cache()
+    route: Final = respx_mock.post(host=GEMINI_HOST, path=GEMINI_GENERATE_CONTENT_PATH).mock(
+        side_effect=[httpx.Response(429, json=RESOURCE_EXHAUSTED)] * CONSECUTIVE_RATE_LIMITS
+        + [httpx.Response(200, json=PONG)]
+    )
+    started: Final = time.monotonic()
+    response: Final = await _router_from_ci_proxy_config().agenerate_content(
+        model="gemini-2.5-flash-lite",
+        contents=[{"role": "user", "parts": [{"text": "Reply with only the single word: pong"}]}],
+    )
+    elapsed: Final = time.monotonic() - started
+
+    assert response.model_dump()["candidates"][0]["content"]["parts"][0]["text"] == "pong"
+    assert route.call_count == CONSECUTIVE_RATE_LIMITS + 1
+    assert elapsed >= MINIMUM_BACKOFF_SECONDS