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.
This commit is contained in:
Yuneng Jiang 2026-09-03 13:49:43 -07:00
parent 16ac67d683
commit 52e24aebba
No known key found for this signature in database
2 changed files with 13 additions and 10 deletions

View file

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

View file

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