mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
test(streaming): cover cancel-race return and str completion_stream
Add test for _put returning when fut.cancel() returns False (line 113) and test for __anext__ str/bytes completion_stream path (line 2193).
This commit is contained in:
parent
4b980d9d9c
commit
8c544a3864
1 changed files with 55 additions and 0 deletions
|
|
@ -2298,6 +2298,61 @@ async def test_queue_wrapper_put_returns_on_loop_exception():
|
|||
wrapper2.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queue_wrapper_put_returns_when_cancel_fails():
|
||||
"""_put returns (no duplicate) when fut.cancel() returns False — item already enqueued (line 112-113)."""
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
mock_future = MagicMock()
|
||||
call_count = {"n": 0}
|
||||
|
||||
def result_side_effect(timeout=None):
|
||||
call_count["n"] += 1
|
||||
if call_count["n"] == 1:
|
||||
raise TimeoutError()
|
||||
return None
|
||||
|
||||
mock_future.result = MagicMock(side_effect=result_side_effect)
|
||||
mock_future.cancel = MagicMock(return_value=False)
|
||||
|
||||
with patch(
|
||||
"asyncio.run_coroutine_threadsafe",
|
||||
return_value=mock_future,
|
||||
):
|
||||
wrapper = _SyncIteratorToQueue(iter(["only-once"]), loop)
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
# fut.cancel() returned False → _put returned without retrying
|
||||
# So run_coroutine_threadsafe should have been called only once for "only-once"
|
||||
assert mock_future.cancel.called
|
||||
wrapper.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anext_str_completion_stream():
|
||||
"""Line 2193: chunk = self.completion_stream when stream is a str/bytes."""
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {"litellm_params": {}}
|
||||
logging_obj.stream_options = None
|
||||
logging_obj.messages = []
|
||||
|
||||
stream = CustomStreamWrapper(
|
||||
completion_stream="raw-string-chunk",
|
||||
model="test-model",
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
# completion_stream is a str → hits line 2193 directly
|
||||
# chunk_creator may raise or return something; we just need the line hit
|
||||
try:
|
||||
await stream.__anext__()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queue_wrapper_cleanup_on_httpx_timeout():
|
||||
"""_queue_wrapper is cleaned up when httpx.TimeoutException is raised during streaming (lines 2306-2307)."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue