async def _stop_flag_checker(
signal_flag: aiotasks.Future,
stop_flag: Optional[primitives.Flag],
) -> None:
"""
A top-level task for external stopping by setting a stop-flag. Once set,
this task will exit, and thus all other top-level tasks will be cancelled.
"""
# Selects the flags to be awaited (if set).
flags = []
if signal_flag is not None:
flags.append(signal_flag)
if stop_flag is not None:
flags.append(aiotasks.create_task(primitives.wait_flag(stop_flag),
name="stop-flag waiter"))
# Wait until one of the stoppers is set/raised.
try:
done, pending = await asyncio.wait(flags, return_when=asyncio.FIRST_COMPLETED)
future = done.pop()
result = await future
except asyncio.CancelledError:
pass # operator is stopping for any other reason
else:
if result is None:
logger.info("Stop-flag is raised. Operator is stopping.")
elif isinstance(result, signal.Signals):
logger.info("Signal %s is received. Operator is stopping.", result.name)
else:
logger.info("Stop-flag is set to %r. Operator is stopping.", result)
async def test_waiting_for_raised_during_the_wait(flag, timer):
async def raise_delayed(delay: float) -> None:
await asyncio.sleep(delay)
await raise_flag(flag) # tested separately above
asyncio.create_task(raise_delayed(0.2))
with timer:
await asyncio.wait_for(wait_flag(flag), timeout=1.0)
assert 0.2 <= timer.seconds < 0.5 # near-instant once raised
async def test_waiting_for_preraised_is_instant(flag, timer):
await raise_flag(flag) # tested separately above
with timer:
await asyncio.wait_for(wait_flag(flag), timeout=1.0)
assert timer.seconds < 0.5 # near-instant, plus code overhead
async def test_waiting_for_unraised_times_out(flag, timer):
with timer:
with pytest.raises(asyncio.TimeoutError):
await asyncio.wait_for(wait_flag(flag), timeout=0.1)
assert timer.seconds >= 0.1
