From 0b9a24202e59db1e23d9f82d958df7e3e3dd20ab Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sun, 15 Mar 2026 17:37:42 -0700 Subject: [PATCH] Fix flaky vertex pass-through spend test by polling instead of fixed sleep The test used a fixed 40s sleep before checking spend logs, but async spend logging in CI sometimes takes longer to flush. Replace with a polling loop (10s interval, 120s max) that exits early on success. Co-Authored-By: Claude Opus 4.6 --- tests/pass_through_tests/test_vertex_ai.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/pass_through_tests/test_vertex_ai.py b/tests/pass_through_tests/test_vertex_ai.py index b3a99bc553..2f5ec8eaa3 100644 --- a/tests/pass_through_tests/test_vertex_ai.py +++ b/tests/pass_through_tests/test_vertex_ai.py @@ -109,17 +109,25 @@ async def test_basic_vertex_ai_pass_through_with_spendlog(): print("response", response) - await asyncio.sleep(40) - spend_after = await call_spend_logs_endpoint() - print("spend_after", spend_after) + # Poll for spend update instead of fixed sleep - spend logging is async/batched + max_wait = 120 # 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 + 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. spend_before: {}, spend_after: {}".format( - spend_before, spend_after + ), "Spend should be greater than before after {}s. spend_before: {}, spend_after: {}".format( + elapsed, spend_before, spend_after ) - pass - @pytest.mark.asyncio() @pytest.mark.skip(reason="skip flaky test - vertex pass through streaming is flaky")