diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/test_litellm/litellm_core_utils/test_logging_worker.py index 336067976fa..2bb93a58531 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -189,23 +189,20 @@ class TestLoggingWorker: ``RuntimeError: ... is bound to a different event loop`` from the queue's Event. """ worker = LoggingWorker(timeout=1.0, max_queue_size=10) - fired = [] - - async def marker(): - fired.append(True) + callback = AsyncMock() async def enqueue_on_first_loop(): if stranded == "still_queued": worker._ensure_queue() - worker.enqueue(marker()) + worker.enqueue(callback()) return - worker.ensure_initialized_and_enqueue(marker()) + worker.ensure_initialized_and_enqueue(callback()) asyncio.run(enqueue_on_first_loop()) assert worker._queue is not None expected_shape = (1, 0) if stranded == "still_queued" else (0, 1) assert (worker._queue.qsize(), len(worker._unstarted_dequeued_tasks())) == expected_shape - assert fired == [], "precondition: the callback never ran before the first loop closed" + assert callback.await_count == 0, "precondition: the callback never ran before the first loop closed" async def flush_twice_on_second_loop(): await asyncio.wait_for(worker.flush(), timeout=5) @@ -213,25 +210,22 @@ class TestLoggingWorker: asyncio.run(flush_twice_on_second_loop()) - assert fired == [True] + assert callback.await_count == 1 def test_flush_starts_a_worker_when_the_queue_has_none(self): """``flush()`` must drain a queue that exists on the current loop without a running worker.""" worker = LoggingWorker(timeout=1.0, max_queue_size=10) - fired = [] - - async def marker(): - fired.append(True) + callback = AsyncMock() async def enqueue_then_flush(): worker._ensure_queue() - worker.enqueue(marker()) + worker.enqueue(callback()) assert worker._worker_task is None, "precondition: nothing is draining the queue yet" await asyncio.wait_for(worker.flush(), timeout=3) asyncio.run(enqueue_then_flush()) - assert fired == [True] + assert callback.await_count == 1 def test_flush_on_exit_swallows_cancellation_and_drains_remaining(self): """A callback raising CancelledError must not abort the atexit flush of later events."""