test_waitfor.py

importasyncio
importunittest
importtime
fromtestimportsupport


deftearDownModule():
asyncio.set_event_loop_policy(None)


# The following value can be used as a very small timeout:
# it passes check "timeout > 0", but has almost
# no effect on the test performance
_EPSILON=0.0001


classSlowTask:
""" Task will run for this defined time, ignoring cancel requests """
TASK_TIMEOUT=0.2

def__init__(self):
self.exited=False

asyncdefrun(self):
exitat=time.monotonic() +self.TASK_TIMEOUT

whileTrue:
tosleep=exitat-time.monotonic()
iftosleep<=0:
break

try:
awaitasyncio.sleep(tosleep)
exceptasyncio.CancelledError:
pass

self.exited=True


classAsyncioWaitForTest(unittest.IsolatedAsyncioTestCase):

asyncdeftest_asyncio_wait_for_cancelled(self):
t=SlowTask()

waitfortask=asyncio.create_task(
asyncio.wait_for(t.run(), t.TASK_TIMEOUT*2))
awaitasyncio.sleep(0)
waitfortask.cancel()
awaitasyncio.wait({waitfortask})

self.assertTrue(t.exited)

asyncdeftest_asyncio_wait_for_timeout(self):
t=SlowTask()

try:
awaitasyncio.wait_for(t.run(), t.TASK_TIMEOUT/2)
exceptasyncio.TimeoutError:
pass

self.assertTrue(t.exited)

asyncdeftest_wait_for_timeout_less_then_0_or_0_future_done(self):
loop=asyncio.get_running_loop()

fut=loop.create_future()
fut.set_result('done')

ret=awaitasyncio.wait_for(fut, 0)

self.assertEqual(ret, 'done')
self.assertTrue(fut.done())

asyncdeftest_wait_for_timeout_less_then_0_or_0_coroutine_do_not_started(self):
foo_started=False

asyncdeffoo():
nonlocalfoo_started
foo_started=True

withself.assertRaises(asyncio.TimeoutError):
awaitasyncio.wait_for(foo(), 0)

self.assertEqual(foo_started, False)

asyncdeftest_wait_for_timeout_less_then_0_or_0(self):
loop=asyncio.get_running_loop()

fortimeoutin [0, -1]:
withself.subTest(timeout=timeout):
foo_running=None
started=loop.create_future()

asyncdeffoo():
nonlocalfoo_running
foo_running=True
started.set_result(None)
try:
awaitasyncio.sleep(10)
finally:
foo_running=False
return'done'

fut=asyncio.create_task(foo())
awaitstarted

withself.assertRaises(asyncio.TimeoutError):
awaitasyncio.wait_for(fut, timeout)

self.assertTrue(fut.done())
# it should have been cancelled due to the timeout
self.assertTrue(fut.cancelled())
self.assertEqual(foo_running, False)

asyncdeftest_wait_for(self):
foo_running=None

asyncdeffoo():
nonlocalfoo_running
foo_running=True
try:
awaitasyncio.sleep(support.LONG_TIMEOUT)
finally:
foo_running=False
return'done'

fut=asyncio.create_task(foo())

withself.assertRaises(asyncio.TimeoutError):
awaitasyncio.wait_for(fut, 0.1)
self.assertTrue(fut.done())
# it should have been cancelled due to the timeout
self.assertTrue(fut.cancelled())
self.assertEqual(foo_running, False)

asyncdeftest_wait_for_blocking(self):
asyncdefcoro():
return'done'

res=awaitasyncio.wait_for(coro(), timeout=None)
self.assertEqual(res, 'done')

asyncdeftest_wait_for_race_condition(self):
loop=asyncio.get_running_loop()

fut=loop.create_future()
task=asyncio.wait_for(fut, timeout=0.2)
loop.call_soon(fut.set_result, "ok")
res=awaittask
self.assertEqual(res, "ok")

asyncdeftest_wait_for_cancellation_race_condition(self):
asyncdefinner():
withself.assertRaises(asyncio.CancelledError):
awaitasyncio.sleep(1)
return1

result=awaitasyncio.wait_for(inner(), timeout=.01)
self.assertEqual(result, 1)

asyncdeftest_wait_for_waits_for_task_cancellation(self):
task_done=False

asyncdefinner():
nonlocaltask_done
try:
awaitasyncio.sleep(10)
exceptasyncio.CancelledError:
awaitasyncio.sleep(_EPSILON)
raise
finally:
task_done=True

inner_task=asyncio.create_task(inner())

withself.assertRaises(asyncio.TimeoutError) ascm:
awaitasyncio.wait_for(inner_task, timeout=_EPSILON)

self.assertTrue(task_done)
chained=cm.exception.__context__
self.assertEqual(type(chained), asyncio.CancelledError)

asyncdeftest_wait_for_waits_for_task_cancellation_w_timeout_0(self):
task_done=False

asyncdeffoo():
asyncdefinner():
nonlocaltask_done
try:
awaitasyncio.sleep(10)
exceptasyncio.CancelledError:
awaitasyncio.sleep(_EPSILON)
raise
finally:
task_done=True

inner_task=asyncio.create_task(inner())
awaitasyncio.sleep(_EPSILON)
awaitasyncio.wait_for(inner_task, timeout=0)

withself.assertRaises(asyncio.TimeoutError) ascm:
awaitfoo()

self.assertTrue(task_done)
chained=cm.exception.__context__
self.assertEqual(type(chained), asyncio.CancelledError)

asyncdeftest_wait_for_reraises_exception_during_cancellation(self):
classFooException(Exception):
pass

asyncdeffoo():
asyncdefinner():
try:
awaitasyncio.sleep(0.2)
finally:
raiseFooException

inner_task=asyncio.create_task(inner())

awaitasyncio.wait_for(inner_task, timeout=_EPSILON)

withself.assertRaises(FooException):
awaitfoo()

asyncdef_test_cancel_wait_for(self, timeout):
loop=asyncio.get_running_loop()

asyncdefblocking_coroutine():
fut=loop.create_future()
# Block: fut result is never set
awaitfut

task=asyncio.create_task(blocking_coroutine())

wait=asyncio.create_task(asyncio.wait_for(task, timeout))
loop.call_soon(wait.cancel)

withself.assertRaises(asyncio.CancelledError):
awaitwait

# Python issue #23219: cancelling the wait must also cancel the task
self.assertTrue(task.cancelled())

asyncdeftest_cancel_blocking_wait_for(self):
awaitself._test_cancel_wait_for(None)

asyncdeftest_cancel_wait_for(self):
awaitself._test_cancel_wait_for(60.0)

asyncdeftest_wait_for_cancel_suppressed(self):
# GH-86296: Suppressing CancelledError is discouraged
# but if a task suppresses CancelledError and returns a value,
# `wait_for` should return the value instead of raising CancelledError.
# This is the same behavior as `asyncio.timeout`.

asyncdefreturn_42():
try:
awaitasyncio.sleep(10)
exceptasyncio.CancelledError:
return42

res=awaitasyncio.wait_for(return_42(), timeout=0.1)
self.assertEqual(res, 42)


asyncdeftest_wait_for_issue86296(self):
# GH-86296: The task should get cancelled and not run to completion.
# inner completes in one cycle of the event loop so it
# completes before the task is cancelled.

asyncdefinner():
return'done'

inner_task=asyncio.create_task(inner())
reached_end=False

asyncdefwait_for_coro():
awaitasyncio.wait_for(inner_task, timeout=100)
awaitasyncio.sleep(1)
nonlocalreached_end
reached_end=True

task=asyncio.create_task(wait_for_coro())
self.assertFalse(task.done())
# Run the task
awaitasyncio.sleep(0)
task.cancel()
withself.assertRaises(asyncio.CancelledError):
awaittask
self.assertTrue(inner_task.done())
self.assertEqual(awaitinner_task, 'done')
self.assertFalse(reached_end)


classWaitForShieldTests(unittest.IsolatedAsyncioTestCase):

asyncdeftest_zero_timeout(self):
# `asyncio.shield` creates a new task which wraps the passed in
# awaitable and shields it from cancellation so with timeout=0
# the task returned by `asyncio.shield` aka shielded_task gets
# cancelled immediately and the task wrapped by it is scheduled
# to run.

asyncdefcoro():
awaitasyncio.sleep(0.01)
return'done'

task=asyncio.create_task(coro())
withself.assertRaises(asyncio.TimeoutError):
shielded_task=asyncio.shield(task)
awaitasyncio.wait_for(shielded_task, timeout=0)

# Task is running in background
self.assertFalse(task.done())
self.assertFalse(task.cancelled())
self.assertTrue(shielded_task.cancelled())

# Wait for the task to complete
awaitasyncio.sleep(0.1)
self.assertTrue(task.done())


asyncdeftest_none_timeout(self):
# With timeout=None the timeout is disabled so it
# runs till completion.
asyncdefcoro():
awaitasyncio.sleep(0.1)
return'done'

task=asyncio.create_task(coro())
awaitasyncio.wait_for(asyncio.shield(task), timeout=None)

self.assertTrue(task.done())
self.assertEqual(awaittask, "done")

asyncdeftest_shielded_timeout(self):
# shield prevents the task from being cancelled.
asyncdefcoro():
awaitasyncio.sleep(0.1)
return'done'

task=asyncio.create_task(coro())
withself.assertRaises(asyncio.TimeoutError):
awaitasyncio.wait_for(asyncio.shield(task), timeout=0.01)

self.assertFalse(task.done())
self.assertFalse(task.cancelled())
self.assertEqual(awaittask, "done")


if__name__=='__main__':
unittest.main()