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 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang 2026-03-15 17:37:42 -07:00
parent 40edb16fb9
commit 0b9a24202e

View File

@ -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")