test(pass-through): de-flake vertex spend-log test by routing through the proxy

tests/pass_through_tests/test_vertex_ai.py

@@ -11,6 +11,7 @@
 import os
 import pytest
 import asyncio
+import requests
 
 # Path to your service account JSON file
 SERVICE_ACCOUNT_FILE = "path/to/your/service-account.json"
@@ -57,106 +58,122 @@ def load_vertex_ai_credentials():
     os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name)
 
 
-async def call_spend_logs_endpoint():
-    """
-    Call this
-    curl -X GET "http://0.0.0.0:4000/spend/logs" -H "Authorization: Bearer sk-1234"
-    """
-    import datetime
-    import requests
+LITE_LLM_ENDPOINT = "http://localhost:4000"
 
-    todays_date = datetime.datetime.now().strftime("%Y-%m-%d")
-    url = f"http://0.0.0.0:4000/global/spend/logs?api_key=best-api-key-ever"
-    headers = {"Authorization": f"Bearer sk-1234"}
-    response = requests.get(url, headers=headers)
-    print("response from call_spend_logs_endpoint", response)
+SPEND_LOG_API_KEY = "best-api-key-ever"
 
-    if response.status_code != 200:
-        print(f"spend logs endpoint returned {response.status_code}: {response.text}")
-        return None
-
-    json_response = response.json()
 
-    # get spend for today
+def get_tracked_spend() -> float:
     """
-    json response looks like this
+    Total spend recorded under the pass-through key in the global spend view.
 
-    [{'date': '2024-08-30', 'spend': 0.00016600000000000002, 'api_key': 'best-api-key-ever'}]
+    Sums every day the endpoint returns instead of matching the runner's local
+    "today" so a UTC date rollover mid-test can't hide a freshly billed call, and
+    treats an unreachable endpoint as "nothing recorded yet" (0.0).
     """
-    print("json_response", json_response)
+    url = f"{LITE_LLM_ENDPOINT}/global/spend/logs?api_key={SPEND_LOG_API_KEY}"
+    response = requests.get(url, headers={"Authorization": "Bearer sk-1234"})
+    if response.status_code != 200:
+        print(f"global spend logs endpoint returned {response.status_code}: {response.text}")
+        return 0.0
 
-    todays_date = datetime.datetime.now().strftime("%Y-%m-%d")
-    for spend_log in json_response:
-        if spend_log["date"] == todays_date:
-            return spend_log["spend"]
+    rows = response.json()
+    print("global spend logs rows", rows)
+    return sum(float(row.get("spend") or 0.0) for row in rows)
 
 
-LITE_LLM_ENDPOINT = "http://localhost:4000"
-
+VERTEX_PROJECT = "litellm-ci-cd"
+VERTEX_MODEL = "gemini-3.1-flash-lite"
+VERTEX_GENERATE_CONTENT_URL = (
+    f"{LITE_LLM_ENDPOINT}/vertex_ai/v1/projects/{VERTEX_PROJECT}"
+    f"/locations/global/publishers/google/models/{VERTEX_MODEL}:generateContent"
+)
 
-def _is_vertex_quota_error(exc: Exception) -> bool:
-    message = str(exc)
-    return (
-        "429" in message
-        or "Too Many Requests" in message
-        or "RESOURCE_EXHAUSTED" in message
-    )
 
+def _vertex_access_token() -> str:
+    import google.auth
+    import google.auth.transport.requests
 
-@pytest.mark.asyncio()
-async def test_basic_vertex_ai_pass_through_with_spendlog():
+    credentials, _ = google.auth.default(
+        scopes=["https://www.googleapis.com/auth/cloud-platform"]
+    )
+    credentials.refresh(google.auth.transport.requests.Request())
+    return credentials.token
 
-    spend_before = await call_spend_logs_endpoint() or 0.0
-    load_vertex_ai_credentials()
 
-    vertexai.init(
-        project="litellm-ci-cd",
-        location="global",
-        api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex_ai",
-        api_transport="rest",
+def _spend_log_for_request(call_id: str) -> dict | None:
+    response = requests.get(
+        f"{LITE_LLM_ENDPOINT}/spend/logs?request_id={call_id}",
+        headers={"Authorization": "Bearer sk-1234"},
+        timeout=30,
     )
+    if response.status_code != 200:
+        return None
+    rows = response.json()
+    return rows[0] if rows else None
 
-    model = GenerativeModel(model_name="gemini-3.1-flash-lite")
-    try:
-        response = model.generate_content("hi")
-    except Exception as exc:
-        if _is_vertex_quota_error(exc):
-            pytest.skip("Vertex AI quota exhausted")
-        raise
 
-    print("response", response)
+def _is_vertex_quota_error(response: requests.Response) -> bool:
+    return response.status_code == 429 or "RESOURCE_EXHAUSTED" in response.text
 
-    # Spend logging is async/batched and can lag under CI load, so poll instead of
-    # sleeping a fixed amount. A transient empty read is skipped, not counted as 0.0
-    # spend, which would spuriously fail the assertion on an otherwise-billed call.
-    max_wait = 240  # total seconds to wait
-    poll_interval = 10  # seconds between checks
-    elapsed = 0
-    spend_after = spend_before
-    while elapsed < max_wait:
-        await asyncio.sleep(poll_interval)
-        elapsed += poll_interval
-        latest_spend = await call_spend_logs_endpoint()
-        if latest_spend is None:
-            print(f"spend logs unavailable (elapsed={elapsed}s), retrying")
-            continue
-        spend_after = latest_spend
-        print(f"spend_after (elapsed={elapsed}s)", spend_after)
-        if spend_after > spend_before:
-            break
 
-    assert (
-        spend_after > spend_before
-    ), "Spend should be greater than before after {}s. spend_before: {}, spend_after: {}".format(
-        elapsed, spend_before, spend_after
+@pytest.mark.asyncio()
+async def test_basic_vertex_ai_pass_through_with_spendlog():
+    load_vertex_ai_credentials()
+    access_token = _vertex_access_token()
+
+    # Drive the pass-through over HTTP instead of the vertexai SDK: the SDK intermittently
+    # routes generateContent to the public Vertex endpoint rather than the proxy override,
+    # so the call never reaches LiteLLM and no spend is logged. A direct request always
+    # hits the proxy. Spend logging then runs on a best-effort background worker that can
+    # drop a single event, so retry a few billed calls and assert that one specific call's
+    # spend log lands. Failing every attempt still fails hard, which is the signal we want
+    # if cost tracking is broken.
+    max_attempts = 3
+    poll_seconds = 60
+    poll_interval = 5
+
+    for attempt in range(1, max_attempts + 1):
+        response = requests.post(
+            VERTEX_GENERATE_CONTENT_URL,
+            headers={
+                "Authorization": f"Bearer {access_token}",
+                "Content-Type": "application/json",
+            },
+            json={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]},
+            timeout=60,
+        )
+        if _is_vertex_quota_error(response):
+            pytest.skip("Vertex AI quota exhausted")
+        assert (
+            response.status_code == 200
+        ), f"vertex pass-through call failed: {response.status_code} {response.text}"
+
+        call_id = response.headers.get("x-litellm-call-id")
+        assert call_id, "proxy response missing x-litellm-call-id header"
+
+        for _ in range(poll_seconds // poll_interval):
+            await asyncio.sleep(poll_interval)
+            row = _spend_log_for_request(call_id)
+            if row is not None and float(row.get("spend") or 0) > 0:
+                assert "gemini" in row["model"], f"unexpected model in spend log: {row}"
+                assert (
+                    row["custom_llm_provider"] == "vertex_ai"
+                ), f"unexpected provider in spend log: {row}"
+                return
+
+        print(f"attempt {attempt}: spend log for call {call_id} not found yet, re-billing")
+
+    pytest.fail(
+        f"Vertex pass-through spend never recorded after {max_attempts} billed calls"
     )
 
 
 @pytest.mark.asyncio()
 @pytest.mark.skip(reason="skip flaky test - vertex pass through streaming is flaky")
 async def test_basic_vertex_ai_pass_through_streaming_with_spendlog():
 
-    spend_before = await call_spend_logs_endpoint() or 0.0
+    spend_before = get_tracked_spend()
     print("spend_before", spend_before)
     load_vertex_ai_credentials()
 
@@ -176,7 +193,7 @@ async def test_basic_vertex_ai_pass_through_streaming_with_spendlog():
     print("response", response)
 
     await asyncio.sleep(20)
-    spend_after = await call_spend_logs_endpoint()
+    spend_after = get_tracked_spend()
     print("spend_after", spend_after)
     assert (
         spend_after > spend_before