mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
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:
parent
16ac67d683
commit
52e24aebba
2 changed files with 13 additions and 10 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue