From c376b8be583ba1a93ab45320803e0808eb70b9fe Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Sat, 28 Feb 2026 08:56:42 +0000
Subject: [PATCH] fix(test): add spend data polling with retries for e2e
 pass-through tests

- test_vertex_with_spend.test.js: Replace 15s fixed wait with polling loop
  (up to 6 attempts, 10s apart) for spend data to appear in DB
- Increase test timeout from 25s to 90s to accommodate polling
- base_anthropic_messages_tool_search_test.py: Add flaky(retries=3) for
  streaming test that depends on live Anthropic API

Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>
---
 .../test_vertex_with_spend.test.js            | 56 +++++++++----------
 ...ase_anthropic_messages_tool_search_test.py |  1 +
 2 files changed, 27 insertions(+), 30 deletions(-)

diff --git a/tests/pass_through_tests/test_vertex_with_spend.test.js b/tests/pass_through_tests/test_vertex_with_spend.test.js
index acfeffd6b9f..1b0291e474f 100644
--- a/tests/pass_through_tests/test_vertex_with_spend.test.js
+++ b/tests/pass_through_tests/test_vertex_with_spend.test.js
@@ -100,21 +100,19 @@ describe('Vertex AI Tests', () => {
         const callId = lastCallId;
         console.log("Captured Call ID:", callId);
 
-        // Wait for spend to be logged
-        await new Promise(resolve => setTimeout(resolve, 15000));
+        // Poll for spend data with retries (DB writes can be slow in CI)
+        let spendData = null;
+        for (let attempt = 0; attempt < 6; attempt++) {
+            await new Promise(resolve => setTimeout(resolve, 10000));
+            const spendResponse = await fetch(
+                `http://127.0.0.1:4000/spend/logs?request_id=${callId}`,
+                { headers: { 'Authorization': 'Bearer sk-1234' } }
+            );
+            spendData = await spendResponse.json();
+            console.log(`spendData (attempt ${attempt + 1}):`, spendData);
+            if (spendData && spendData.length > 0 && spendData[0] && spendData[0].request_id) break;
+        }
 
-        // Check spend logs
-        const spendResponse = await fetch(
-            `http://127.0.0.1:4000/spend/logs?request_id=${callId}`,
-            {
-                headers: {
-                    'Authorization': 'Bearer sk-1234'
-                }
-            }
-        );
-        
-        const spendData = await spendResponse.json();
-        console.log("spendData", spendData)
         expect(spendData).toBeDefined();
         expect(spendData[0].request_id).toBe(callId);
         expect(spendData[0].call_type).toBe('pass_through_endpoint');
@@ -123,7 +121,7 @@ describe('Vertex AI Tests', () => {
         expect(spendData[0].model).toContain('gemini');
         expect(spendData[0].spend).toBeGreaterThan(0);
         expect(spendData[0].custom_llm_provider).toBe('vertex_ai');
-    }, 25000);
+    }, 90000);
 
     test('should successfully generate streaming content with tags', async () => {
         const vertexAI = new VertexAI({
@@ -170,21 +168,19 @@ describe('Vertex AI Tests', () => {
         const callId = lastCallId;
         console.log("Captured Call ID:", callId);
 
-        // Wait for spend to be logged
-        await new Promise(resolve => setTimeout(resolve, 15000));
+        // Poll for spend data with retries (DB writes can be slow in CI)
+        let spendData = null;
+        for (let attempt = 0; attempt < 6; attempt++) {
+            await new Promise(resolve => setTimeout(resolve, 10000));
+            const spendResponse = await fetch(
+                `http://127.0.0.1:4000/spend/logs?request_id=${callId}`,
+                { headers: { 'Authorization': 'Bearer sk-1234' } }
+            );
+            spendData = await spendResponse.json();
+            console.log(`spendData (attempt ${attempt + 1}):`, spendData);
+            if (spendData && spendData.length > 0 && spendData[0] && spendData[0].request_id) break;
+        }
 
-        // Check spend logs
-        const spendResponse = await fetch(
-            `http://127.0.0.1:4000/spend/logs?request_id=${callId}`,
-            {
-                headers: {
-                    'Authorization': 'Bearer sk-1234'
-                }
-            }
-        );
-        
-        const spendData = await spendResponse.json();
-        console.log("spendData", spendData)
         expect(spendData).toBeDefined();
         expect(spendData[0].request_id).toBe(callId);
         expect(spendData[0].call_type).toBe('pass_through_endpoint');
@@ -193,5 +189,5 @@ describe('Vertex AI Tests', () => {
         expect(spendData[0].model).toContain('gemini');
         expect(spendData[0].spend).toBeGreaterThan(0);
         expect(spendData[0].custom_llm_provider).toBe('vertex_ai');
-    }, 25000);
+    }, 90000);
 });
\ No newline at end of file
diff --git a/tests/pass_through_unit_tests/base_anthropic_messages_tool_search_test.py b/tests/pass_through_unit_tests/base_anthropic_messages_tool_search_test.py
index 590e746b39c..045c43c1ff8 100644
--- a/tests/pass_through_unit_tests/base_anthropic_messages_tool_search_test.py
+++ b/tests/pass_through_unit_tests/base_anthropic_messages_tool_search_test.py
@@ -206,6 +206,7 @@ class BaseAnthropicMessagesToolSearchTest(ABC):
             assert len(tool_uses) > 0, "Expected tool_use blocks when stop_reason is tool_use"
 
     @pytest.mark.asyncio
+    @pytest.mark.flaky(retries=3, delay=5)
     async def test_tool_search_streaming(self):
         """
         E2E test: Tool search should work with streaming responses.