From b3f04e733081b134ce0a886f00ccc3c66e81bfbb Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 17 Jun 2026 20:20:00 +0000
Subject: [PATCH] test(pass_through): harden vertex spendlog poll against
 transient empty reads

test_basic_vertex_ai_pass_through_with_spendlog failed intermittently on
litellm_internal_staging (pipelines 82155, 82196, 82209, 82230) with "Spend
should be greater than before after 120s". Spend logging is async and batched,
so the pass-through call's cost sometimes had not landed within the 120s poll
window; one run ended on spend_after 0.0 because the final /global/spend/logs
read returned nothing and "or 0.0" recorded that as zero spend.

Widen the poll window to 240s and skip a transient empty read instead of
treating it as 0.0, so a momentary endpoint hiccup on the last poll no longer
fails an otherwise-billed call. The spend_after > spend_before assertion is
unchanged, so a genuinely unbilled call still fails the test
---
 tests/pass_through_tests/test_vertex_ai.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/tests/pass_through_tests/test_vertex_ai.py b/tests/pass_through_tests/test_vertex_ai.py
index 0ac66b470c6f..e8223f2219c7 100644
--- a/tests/pass_through_tests/test_vertex_ai.py
+++ b/tests/pass_through_tests/test_vertex_ai.py
@@ -126,15 +126,21 @@ async def test_basic_vertex_ai_pass_through_with_spendlog():
 
     print("response", response)
 
-    # Poll for spend update instead of fixed sleep - spend logging is async/batched
-    max_wait = 120  # total seconds to wait
+    # 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
-        spend_after = await call_spend_logs_endpoint() or 0.0
+        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