diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index 227d2966e53..c8d05bc0f57 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -4,10 +4,10 @@ Tests for AnthropicResponsesStreamWrapper """ import asyncio +import datetime import json import os import sys -import time from collections.abc import Awaitable, Callable from types import SimpleNamespace @@ -641,7 +641,7 @@ def _logging_obj_with_failure_hook(on_failure: Callable[[Exception], Awaitable[N messages=[{"role": "user", "content": "hi"}], stream=True, call_type="anthropic_messages", - start_time=time.time(), + start_time=datetime.datetime(2026, 1, 1), litellm_call_id="lit-6871-test", function_id="lit-6871-test", ) diff --git a/tests/unit/responses/test_responses_streaming_iterator.py b/tests/unit/responses/test_responses_streaming_iterator.py index 16b125e56d4..02ce972268f 100644 --- a/tests/unit/responses/test_responses_streaming_iterator.py +++ b/tests/unit/responses/test_responses_streaming_iterator.py @@ -174,6 +174,15 @@ class FailureRecorder(CustomLogger): self.success_payloads = (*self.success_payloads, kwargs["standard_logging_object"]) +class _ChunkRecorder: + def __init__(self) -> None: + self.chunks: tuple = () + + async def drain(self, iterator: ResponsesAPIStreamingIterator) -> None: + async for chunk in iterator: + self.chunks = (*self.chunks, chunk) + + class _UpstreamThatTimesOutAfterThreeDeltas: headers: dict = {} @@ -211,20 +220,13 @@ async def test_mid_stream_read_timeout_logs_failure_with_partial_usage_and_no_su logging_obj=logging_obj, custom_llm_provider="openai", ) - received: list = [] - - async def _drain(): - async for chunk in iterator: - received.append(chunk) + received = _ChunkRecorder() with pytest.raises(httpx.ReadTimeout, match="Timeout on reading data from socket"): - await _drain() - assert [chunk.delta for chunk in received] == ["The sea ", "is wide ", "and deep."] + await received.drain(iterator) + assert [chunk.delta for chunk in received.chunks] == ["The sea ", "is wide ", "and deep."] - for _ in range(50): - if recorder.failure_payloads: - break - await asyncio.sleep(0.1) + await iterator._await_pending_logging() (payload,) = recorder.failure_payloads assert payload["status"] == "failure" assert payload["prompt_tokens"] > 0