def test_task_cancel_awaiting():
task_mgr = core.AsyncTaskManager.get_global_ptr()
task_chain = task_mgr.make_task_chain("test_task_cancel_awaiting")
fut = core.AsyncFuture()
async def task_main(task):
await fut
return task.done
task = core.PythonTask(task_main, 'task_main')
task.set_task_chain(task_chain.name)
task_mgr.add(task)
task_chain.poll()
assert not task.done()
task_chain.poll()
assert not task.done()
task.cancel()
task_chain.poll()
assert task.done()
assert task.cancelled()
assert fut.done()
assert fut.cancelled()
def test_task_cancel_waiting():
# Calling result() in a threaded task chain should cancel the future being
# waited on if the surrounding task is cancelled.
task_mgr = core.AsyncTaskManager.get_global_ptr()
task_chain = task_mgr.make_task_chain("test_task_cancel_waiting")
task_chain.set_num_threads(1)
fut = core.AsyncFuture()
async def task_main(task):
# This will block the thread this task is in until the future is done,
# or until the task is cancelled (which implicitly cancels the future).
fut.result()
return task.done
task = core.PythonTask(task_main, 'task_main')
task.set_task_chain(task_chain.name)
task_mgr.add(task)
task_chain.start_threads()
try:
assert not task.done()
fut.cancel()
task.wait()
assert task.cancelled()
assert fut.cancelled()
finally:
task_chain.stop_threads()
def test_task_cancel():
task_mgr = core.AsyncTaskManager.get_global_ptr()
task = core.PythonTask(lambda task: task.done)
task_mgr.add(task)
assert not task.done()
task_mgr.remove(task)
assert task.done()
assert task.cancelled()
with pytest.raises(CancelledError):
task.result()
def test_coro_exception():
task_mgr = core.AsyncTaskManager.get_global_ptr()
task_chain = task_mgr.make_task_chain("test_coro_exception")
def coro_main():
raise RuntimeError
yield None
task = core.PythonTask(coro_main())
task.set_task_chain(task_chain.name)
task_mgr.add(task)
task_chain.wait_for_tasks()
assert task.done()
assert not task.cancelled()
with pytest.raises(RuntimeError):
task.result()
def test_task_result():
task_mgr = core.AsyncTaskManager.get_global_ptr()
task_chain = task_mgr.make_task_chain("test_task_result")
def task_main(task):
task.set_result(42)
# It won't yet be marked done until after it returns.
assert not task.done()
return core.PythonTask.done
task = core.PythonTask(task_main)
task.set_task_chain(task_chain.name)
task_mgr.add(task)
task_chain.wait_for_tasks()
assert task.done()
assert not task.cancelled()
assert task.result() == 42
def test_task_cancel_during_run():
task_mgr = core.AsyncTaskManager.get_global_ptr()
task_chain = task_mgr.make_task_chain("test_task_cancel_during_run")
def task_main(task):
task.remove()
# It won't yet be marked done until after it returns.
assert not task.done()
return task.done
task = core.PythonTask(task_main)
task.set_task_chain(task_chain.name)
task_mgr.add(task)
task_chain.wait_for_tasks()
assert task.done()
assert task.cancelled()
with pytest.raises(CancelledError):
task.result()
def test_coro_await_external_cancel_inner():
# Cancel external future being awaited by a coro.
fut = MockFuture()
async def coro_main():
await fut
task = core.PythonTask(coro_main(), 'coro_main')
task_mgr = core.AsyncTaskManager.get_global_ptr()
task_mgr.add(task)
for i in range(5):
task_mgr.poll()
assert not task.done()
fut._state = 'CANCELLED'
assert not task.done()
task_mgr.poll()
assert task.done()
assert task.cancelled()
def test_coro_await_external_cancel_outer():
# Cancel task that is awaiting external future.
fut = MockFuture()
result = []
async def coro_main():
result.append(await fut)
task = core.PythonTask(coro_main(), 'coro_main')
task_mgr = core.AsyncTaskManager.get_global_ptr()
task_mgr.add(task)
for i in range(5):
task_mgr.poll()
assert not task.done()
fut._state = 'CANCELLED'
assert not task.done()
task_mgr.poll()
assert task.done()
assert task.cancelled()
def test_coro_await_cancel_resistant_coro():
# Await another coro in a coro, but cancel the outer.
fut = core.AsyncFuture()
cancelled_caught = [0]
keep_going = [False]
async def cancel_resistant_coro():
while not fut.done():
try:
await core.AsyncFuture.shield(fut)
except CancelledError as ex:
cancelled_caught[0] += 1
async def coro_main():
await cancel_resistant_coro()
task = core.PythonTask(coro_main(), 'coro_main')
task_mgr = core.AsyncTaskManager.get_global_ptr()
task_mgr.add(task)
assert not task.done()
task_mgr.poll()
assert not task.done()
# No cancelling it once it started...
for i in range(3):
assert task.cancel()
assert not task.done()
for j in range(3):
task_mgr.poll()
assert not task.done()
assert cancelled_caught[0] == 3
fut.set_result(None)
task_mgr.poll()
assert task.done()
assert not task.cancelled()
def test_coro_await_external():
# Await an external future in a coro.
fut = MockFuture()
fut._result = 12345
res = []
async def coro_main():
res.append(await fut)
task = core.PythonTask(coro_main(), 'coro_main')
task_mgr = core.AsyncTaskManager.get_global_ptr()
task_mgr.add(task)
for i in range(5):
task_mgr.poll()
assert not task.done()
fut._state = 'FINISHED'
task_mgr.poll()
assert task.done()
assert not task.cancelled()
assert res == [12345]
def test_coro_await_coro():
# Await another coro in a coro.
fut = core.AsyncFuture()
async def coro2():
await fut
async def coro_main():
await coro2()
task = core.PythonTask(coro_main())
task_mgr = core.AsyncTaskManager.get_global_ptr()
task_mgr.add(task)
for i in range(5):
task_mgr.poll()
assert not task.done()
fut.set_result(None)
task_mgr.poll()
assert task.done()
assert not task.cancelled()
