async def test_empty_scheduler_lifecycle(timer):
with timer, async_timeout.timeout(1):
scheduler = Scheduler()
assert scheduler.empty()
await scheduler.wait()
assert scheduler.empty()
await scheduler.close()
assert scheduler.empty()
assert timer.seconds < CODE_OVERHEAD
async def _wait_for_depletion(
*,
signaller: asyncio.Condition,
scheduler: aiotasks.Scheduler,
settings: configuration.OperatorSettings,
streams: Streams,
) -> None:
# Notify all the workers to finish now. Wake them up if they are waiting in the queue-getting.
for stream in streams.values():
await stream.backlog.put(EOS.token)
# Wait for the queues to be depleted, but only if there are some workers running.
# Continue with the tasks termination if the timeout is reached, no matter the queues.
# NB: the scheduler is checked for a case of mocked workers; otherwise, the streams are enough.
async with signaller:
try:
await asyncio.wait_for(
signaller.wait_for(lambda: not streams or scheduler.empty()),
timeout=settings.batching.exit_timeout)
except asyncio.TimeoutError:
pass
# The last check if the termination is going to be graceful or not.
if streams:
logger.warning(
f"Unprocessed streams left for {list(streams.keys())!r}.")
async def test_tasks_are_pending_if_limit_is_reached(timer):
"""
time: ////////----------------------0.1s------------------0.2s--///
task1: ->spawn->start->sleep->finish->|
task2: ->spawn->.....(pending)......->start->sleep->finish->|
"""
task1_started = asyncio.Event()
task1_finished = asyncio.Event()
task2_started = asyncio.Event()
task2_finished = asyncio.Event()
scheduler = Scheduler(limit=1)
with timer, async_timeout.timeout(1):
await scheduler.spawn(f(Mock(), task1_started, 0.1, task1_finished))
await scheduler.spawn(f(Mock(), task2_started, 0.1, task2_finished))
assert timer.seconds < CODE_OVERHEAD # i.e. spawning is not not blocking
with timer, async_timeout.timeout(1):
await task1_finished.wait()
assert not task2_started.is_set()
await task2_finished.wait()
assert timer.seconds > 0.2
assert timer.seconds < 0.2 + CODE_OVERHEAD * 2
await scheduler.close()
async def test_successes_are_not_reported():
exception_handler = Mock()
scheduler = Scheduler(exception_handler=exception_handler)
await scheduler.spawn(f(Mock()))
await scheduler.wait()
await scheduler.close()
assert exception_handler.call_count == 0
async def test_tasks_are_parallel_if_limit_is_not_reached(timer):
"""
time: ////////----------------------0.1s------------------0.2s--///
task1: ->spawn->start->sleep->finish->|
task2: ->spawn->start->sleep->finish->|
"""
task1_started = asyncio.Event()
task1_finished = asyncio.Event()
task2_started = asyncio.Event()
task2_finished = asyncio.Event()
scheduler = Scheduler(limit=2)
with timer:
await scheduler.spawn(f(Mock(), task1_started, 0.1, task1_finished))
await scheduler.spawn(f(Mock(), task2_started, 0.1, task2_finished))
assert timer.seconds < CODE_OVERHEAD # i.e. spawning is not not blocking
with timer:
await task1_finished.wait()
assert task2_started.is_set()
await task2_finished.wait()
assert timer.seconds > 0.1
assert timer.seconds < 0.1 + CODE_OVERHEAD
await scheduler.close()
async def test_cancellations_are_not_reported():
exception_handler = Mock()
mock = Mock(side_effect=asyncio.CancelledError())
scheduler = Scheduler(exception_handler=exception_handler)
await scheduler.spawn(f(mock, 1))
await scheduler.wait()
await scheduler.close()
assert exception_handler.call_count == 0
async def test_no_tasks_are_accepted_after_closing():
scheduler = Scheduler()
await scheduler.close()
assert scheduler._closed
assert scheduler._spawning_task.done()
assert scheduler._cleaning_task.done()
with pytest.raises(RuntimeError, match=r"Cannot add new coroutines"):
await scheduler.spawn(f(Mock(), 1.0))
async def test_exceptions_are_reported():
exception = ValueError('hello')
exception_handler = Mock()
mock = Mock(side_effect=exception)
scheduler = Scheduler(exception_handler=exception_handler)
await scheduler.spawn(f(mock))
await scheduler.wait()
await scheduler.close()
assert exception_handler.call_count == 1
assert exception_handler.call_args[0][0] is exception
async def test_task_spawning_and_cancellation(timer):
mock = Mock()
flag1 = asyncio.Event()
flag2 = asyncio.Event()
scheduler = Scheduler()
result = await scheduler.spawn(f(mock, flag1, 1.0, flag2))
assert result is None
with timer, async_timeout.timeout(1):
await flag1.wait()
assert timer.seconds < CODE_OVERHEAD
assert mock.call_args[0][0] == 'started'
with timer, async_timeout.timeout(1):
await scheduler.close()
assert timer.seconds < CODE_OVERHEAD # near-instant
assert mock.call_args[0][0] == 'cancelled'
async def test_task_spawning_and_graceful_finishing(timer):
mock = Mock()
flag1 = asyncio.Event()
flag2 = asyncio.Event()
scheduler = Scheduler()
result = await scheduler.spawn(f(mock, flag1, 0.1, flag2))
assert result is None
with timer, async_timeout.timeout(1):
await flag1.wait()
assert timer.seconds < CODE_OVERHEAD
assert mock.call_args[0][0] == 'started'
with timer, async_timeout.timeout(1):
await flag2.wait()
assert timer.seconds > 0.1
assert timer.seconds < 0.1 + CODE_OVERHEAD
assert mock.call_args[0][0] == 'finished'
await scheduler.close()