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