From 52e24aebbabdd4d889dda96f3ebeab0e3bd8c3b1 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 3 Sep 2026 13:49:43 -0700
Subject: [PATCH] refactor(tests): assign the streamed id and lock poll once
 instead of rebinding

The cancel test accumulated chunk_count and reassigned response_id on every
iteration, and the lock watcher rebound its query result on every poll. Both are
the mutable-local pattern the repo avoids.

The stream now drains through a generator that stops at the first chunk carrying
a response id, so the caller binds streamed_ids once and reads the id off the
tail. Empty stream, no-id stream and first-chunk-id all behave exactly as the
loop did. The watcher inlines its poll result.
---
 .../test_e2e_openai_responses_api.py          | 20 +++++++++++--------
 .../test_team_delete_member_add_race.py       |  3 +--
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py
index b24ac0bdb96..755f17c394e 100644
--- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py
+++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py
@@ -153,6 +153,15 @@ def test_cancel_response():
             raise e
 
 
+def _response_ids_until_first(stream):
+    """Yield each streamed chunk's response id, stopping at the first chunk that carries one."""
+    for chunk in stream:
+        response_id = getattr(getattr(chunk, "response", None), "id", None)
+        yield response_id
+        if response_id is not None:
+            return
+
+
 def test_cancel_streaming_response():
     """Cancel a background streaming response while it is still generating.
 
@@ -169,15 +178,10 @@ def test_cancel_streaming_response():
         stream=True,
         background=True,
     ) as stream:
-        chunk_count = 0
-        response_id = None
-        for chunk in stream:
-            chunk_count += 1
-            response_id = getattr(getattr(chunk, "response", None), "id", None)
-            if response_id is not None:
-                break
+        streamed_ids = tuple(_response_ids_until_first(stream))
 
-    assert chunk_count > 0, "stream produced no chunks"
+    assert streamed_ids, "stream produced no chunks"
+    response_id = streamed_ids[-1]
     assert response_id is not None, "no streamed chunk carried a response id to cancel"
 
     cancel_response = client.responses.cancel(response_id)
diff --git a/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py
index 9af742b0dfe..d3ffcf2445e 100644
--- a/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py
+++ b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py
@@ -78,8 +78,7 @@ async def _await_lock_contention(watcher, lock_key: tuple[int, int], task, what:
     while time.monotonic() < deadline:
         if task.done():
             raise AssertionError(f"{what} returned without waiting on the team's advisory lock") from task.exception()
-        rows = await watcher.query_raw(_LOCK_WAITER_SQL, classid, objid)
-        if rows[0]["waiters"]:
+        if (await watcher.query_raw(_LOCK_WAITER_SQL, classid, objid))[0]["waiters"]:
             return
         await asyncio.sleep(_LOCK_POLL_SECONDS)
     raise AssertionError(f"{what} never queued on the team's advisory lock within {_LOCK_WAIT_TIMEOUT_SECONDS}s")