# Adapted with permission from the EdgeDB project; # license: PSFL. import sys import gc import asyncio import contextvars import contextlib from asyncio import taskgroups import unittest import warnings from test.test_asyncio.utils import await_without_task # To prevent a warning "test altered the execution environment" def tearDownModule(): asyncio.set_event_loop_policy(None) class MyExc(Exception): pass class MyBaseExc(BaseException): pass def get_error_types(eg): return {type(exc) for exc in eg.exceptions} def no_other_refs(): # due to gh-124392 coroutines now refer to their locals coro = asyncio.current_task().get_coro() frame = sys._getframe(1) while coro.cr_frame != frame: coro = coro.cr_await return [coro] class TestTaskGroup(unittest.IsolatedAsyncioTestCase): async def test_taskgroup_01(self): async def foo1(): await asyncio.sleep(0.1) return 42 async def foo2(): await asyncio.sleep(0.2) return 11 async with taskgroups.TaskGroup() as g: t1 = g.create_task(foo1()) t2 = g.create_task(foo2()) self.assertEqual(t1.result(), 42) self.assertEqual(t2.result(), 11) async def test_taskgroup_02(self): async def foo1(): await asyncio.sleep(0.1) return 42 async def foo2(): await asyncio.sleep(0.2) return 11 async with taskgroups.TaskGroup() as g: t1 = g.create_task(foo1()) await asyncio.sleep(0.15) t2 = g.create_task(foo2()) self.assertEqual(t1.result(), 42) self.assertEqual(t2.result(), 11) async def test_taskgroup_03(self): async def foo1(): await asyncio.sleep(1) return 42 async def foo2(): await asyncio.sleep(0.2) return 11 async with taskgroups.TaskGroup() as g: t1 = g.create_task(foo1()) await asyncio.sleep(0.15) # cancel t1 explicitly, i.e. everything should continue # working as expected. t1.cancel() t2 = g.create_task(foo2()) self.assertTrue(t1.cancelled()) self.assertEqual(t2.result(), 11) async def test_taskgroup_04(self): NUM = 0 t2_cancel = False t2 = None async def foo1(): await asyncio.sleep(0.1) 1 / 0 async def foo2(): nonlocal NUM, t2_cancel try: await asyncio.sleep(1) except asyncio.CancelledError: t2_cancel = True raise NUM += 1 async def runner(): nonlocal NUM, t2 async with taskgroups.TaskGroup() as g: g.create_task(foo1()) t2 = g.create_task(foo2()) NUM += 10 with self.assertRaises(ExceptionGroup) as cm: await asyncio.create_task(runner()) self.assertEqual(get_error_types(cm.exception), {ZeroDivisionError}) self.assertEqual(NUM, 0) self.assertTrue(t2_cancel) self.assertTrue(t2.cancelled()) async def test_cancel_children_on_child_error(self): # When a child task raises an error, the rest of the children # are cancelled and the errors are gathered into an EG. NUM = 0 t2_cancel = False runner_cancel = False async def foo1(): await asyncio.sleep(0.1) 1 / 0 async def foo2(): nonlocal NUM, t2_cancel try: await asyncio.sleep(5) except asyncio.CancelledError: t2_cancel = True raise NUM += 1 async def runner(): nonlocal NUM, runner_cancel async with taskgroups.TaskGroup() as g: g.create_task(foo1()) g.create_task(foo1()) g.create_task(foo1()) g.create_task(foo2()) try: await asyncio.sleep(10) except asyncio.CancelledError: runner_cancel = True raise NUM += 10 # The 3 foo1 sub tasks can be racy when the host is busy - if the # cancellation happens in the middle, we'll see partial sub errors here with self.assertRaises(ExceptionGroup) as cm: await asyncio.create_task(runner()) self.assertEqual(get_error_types(cm.exception), {ZeroDivisionError}) self.assertEqual(NUM, 0) self.assertTrue(t2_cancel) self.assertTrue(runner_cancel) async def test_cancellation(self): NUM = 0 async def foo(): nonlocal NUM try: await asyncio.sleep(5) except asyncio.CancelledError: NUM += 1 raise async def runner(): async with taskgroups.TaskGroup() as g: for _ in range(5): g.create_task(foo()) r = asyncio.create_task(runner()) await asyncio.sleep(0.1) self.assertFalse(r.done()) r.cancel() with self.assertRaises(asyncio.CancelledError) as cm: await r self.assertEqual(NUM, 5) async def test_taskgroup_07(self): NUM = 0 async def foo(): nonlocal NUM try: await asyncio.sleep(5) except asyncio.CancelledError: NUM += 1 raise async def runner(): nonlocal NUM async with taskgroups.TaskGroup() as g: for _ in range(5): g.create_task(foo()) try: await asyncio.sleep(10) except asyncio.CancelledError: NUM += 10 raise r = asyncio.create_task(runner()) await asyncio.sleep(0.1) self.assertFalse(r.done()) r.cancel() with self.assertRaises(asyncio.CancelledError): await r self.assertEqual(NUM, 15) async def test_taskgroup_08(self): async def foo(): try: await asyncio.sleep(10) finally: 1 / 0 async def runner(): async with taskgroups.TaskGroup() as g: for _ in range(5): g.create_task(foo()) await asyncio.sleep(10) r = asyncio.create_task(runner()) await asyncio.sleep(0.1) self.assertFalse(r.done()) r.cancel() with self.assertRaises(ExceptionGroup) as cm: await r self.assertEqual(get_error_types(cm.exception), {ZeroDivisionError}) async def test_taskgroup_09(self): t1 = t2 = None async def foo1(): await asyncio.sleep(1) return 42 async def foo2(): await asyncio.sleep(2) return 11 async def runner(): nonlocal t1, t2 async with taskgroups.TaskGroup() as g: t1 = g.create_task(foo1()) t2 = g.create_task(foo2()) await asyncio.sleep(0.1) 1 / 0 try: await runner() except ExceptionGroup as t: self.assertEqual(get_error_types(t), {ZeroDivisionError}) else: self.fail('ExceptionGroup was not raised') self.assertTrue(t1.cancelled()) self.assertTrue(t2.cancelled()) async def test_taskgroup_10(self): t1 = t2 = None async def foo1(): await asyncio.sleep(1) return 42 async def foo2(): await asyncio.sleep(2) return 11 async def runner(): nonlocal t1, t2 async with taskgroups.TaskGroup() as g: t1 = g.create_task(foo1()) t2 = g.create_task(foo2()) 1 / 0 try: await runner() except ExceptionGroup as t: self.assertEqual(get_error_types(t), {ZeroDivisionError}) else: self.fail('ExceptionGroup was not raised') self.assertTrue(t1.cancelled()) self.assertTrue(t2.cancelled()) async def test_taskgroup_11(self): async def foo(): try: await asyncio.sleep(10) finally: 1 / 0 async def runner(): async with taskgroups.TaskGroup(): async with taskgroups.TaskGroup() as g2: for _ in range(5): g2.create_task(foo()) await asyncio.sleep(10) r = asyncio.create_task(runner()) await asyncio.sleep(0.1) self.assertFalse(r.done()) r.cancel() with self.assertRaises(ExceptionGroup) as cm: await r self.assertEqual(get_error_types(cm.exception), {ExceptionGroup}) self.assertEqual(get_error_types(cm.exception.exceptions[0]), {ZeroDivisionError}) async def test_taskgroup_12(self): async def foo(): try: await asyncio.sleep(10) finally: 1 / 0 async def runner(): async with taskgroups.TaskGroup() as g1: g1.create_task(asyncio.sleep(10)) async with taskgroups.TaskGroup() as g2: for _ in range(5): g2.create_task(foo()) await asyncio.sleep(10) r = asyncio.create_task(runner()) await asyncio.sleep(0.1) self.assertFalse(r.done()) r.cancel() with self.assertRaises(ExceptionGroup) as cm: await r self.assertEqual(get_error_types(cm.exception), {ExceptionGroup}) self.assertEqual(get_error_types(cm.exception.exceptions[0]), {ZeroDivisionError}) async def test_taskgroup_13(self): async def crash_after(t): await asyncio.sleep(t) raise ValueError(t) async def runner(): async with taskgroups.TaskGroup() as g1: g1.create_task(crash_after(0.1)) async with taskgroups.TaskGroup() as g2: g2.create_task(crash_after(10)) r = asyncio.create_task(runner()) with self.assertRaises(ExceptionGroup) as cm: await r self.assertEqual(get_error_types(cm.exception), {ValueError}) async def test_taskgroup_14(self): async def crash_after(t): await asyncio.sleep(t) raise ValueError(t) async def runner(): async with taskgroups.TaskGroup() as g1: g1.create_task(crash_after(10)) async with taskgroups.TaskGroup() as g2: g2.create_task(crash_after(0.1)) r = asyncio.create_task(runner()) with self.assertRaises(ExceptionGroup) as cm: await r self.assertEqual(get_error_types(cm.exception), {ExceptionGroup}) self.assertEqual(get_error_types(cm.exception.exceptions[0]), {ValueError}) async def test_taskgroup_15(self): async def crash_soon(): await asyncio.sleep(0.3) 1 / 0 async def runner(): async with taskgroups.TaskGroup() as g1: g1.create_task(crash_soon()) try: await asyncio.sleep(10) except asyncio.CancelledError: await asyncio.sleep(0.5) raise r = asyncio.create_task(runner()) await asyncio.sleep(0.1) self.assertFalse(r.done()) r.cancel() with self.assertRaises(ExceptionGroup) as cm: await r self.assertEqual(get_error_types(cm.exception), {ZeroDivisionError}) async def test_taskgroup_16(self): async def crash_soon(): await asyncio.sleep(0.3) 1 / 0 async def nested_runner(): async with taskgroups.TaskGroup() as g1: g1.create_task(crash_soon()) try: await asyncio.sleep(10) except asyncio.CancelledError: await asyncio.sleep(0.5) raise async def runner(): t = asyncio.create_task(nested_runner()) await t r = asyncio.create_task(runner()) await asyncio.sleep(0.1) self.assertFalse(r.done()) r.cancel() with self.assertRaises(ExceptionGroup) as cm: await r self.assertEqual(get_error_types(cm.exception), {ZeroDivisionError}) async def test_taskgroup_17(self): NUM = 0 async def runner(): nonlocal NUM async with taskgroups.TaskGroup(): try: await asyncio.sleep(10) except asyncio.CancelledError: NUM += 10 raise r = asyncio.create_task(runner()) await asyncio.sleep(0.1) self.assertFalse(r.done()) r.cancel() with self.assertRaises(asyncio.CancelledError): await r self.assertEqual(NUM, 10) async def test_taskgroup_18(self): NUM = 0 async def runner(): nonlocal NUM async with taskgroups.TaskGroup(): try: await asyncio.sleep(10) except asyncio.CancelledError: NUM += 10 # This isn't a good idea, but we have to support # this weird case. raise MyExc r = asyncio.create_task(runner()) await asyncio.sleep(0.1) self.assertFalse(r.done()) r.cancel() try: await r except ExceptionGroup as t: self.assertEqual(get_error_types(t),{MyExc}) else: self.fail('ExceptionGroup was not raised') self.assertEqual(NUM, 10) async def test_taskgroup_19(self): async def crash_soon(): await asyncio.sleep(0.1) 1 / 0 async def nested(): try: await asyncio.sleep(10) finally: raise MyExc async def runner(): async with taskgroups.TaskGroup() as g: g.create_task(crash_soon()) await nested() r = asyncio.create_task(runner()) try: await r except ExceptionGroup as t: self.assertEqual(get_error_types(t), {MyExc, ZeroDivisionError}) else: self.fail('TasgGroupError was not raised') async def test_taskgroup_20(self): async def crash_soon(): await asyncio.sleep(0.1) 1 / 0 async def nested(): try: await asyncio.sleep(10) finally: raise KeyboardInterrupt async def runner(): async with taskgroups.TaskGroup() as g: g.create_task(crash_soon()) await nested() with self.assertRaises(KeyboardInterrupt): await runner() async def test_taskgroup_20a(self): async def crash_soon(): await asyncio.sleep(0.1) 1 / 0 async def nested(): try: await asyncio.sleep(10) finally: raise MyBaseExc async def runner(): async with taskgroups.TaskGroup() as g: g.create_task(crash_soon()) await nested() with self.assertRaises(BaseExceptionGroup) as cm: await runner() self.assertEqual( get_error_types(cm.exception), {MyBaseExc, ZeroDivisionError} ) async def _test_taskgroup_21(self): # This test doesn't work as asyncio, currently, doesn't # correctly propagate KeyboardInterrupt (or SystemExit) -- # those cause the event loop itself to crash. # (Compare to the previous (passing) test -- that one raises # a plain exception but raises KeyboardInterrupt in nested(); # this test does it the other way around.) async def crash_soon(): await asyncio.sleep(0.1) raise KeyboardInterrupt async def nested(): try: await asyncio.sleep(10) finally: raise TypeError async def runner(): async with taskgroups.TaskGroup() as g: g.create_task(crash_soon()) await nested() with self.assertRaises(KeyboardInterrupt): await runner() async def test_taskgroup_21a(self): async def crash_soon(): await asyncio.sleep(0.1) raise MyBaseExc async def nested(): try: await asyncio.sleep(10) finally: raise TypeError async def runner(): async with taskgroups.TaskGroup() as g: g.create_task(crash_soon()) await nested() with self.assertRaises(BaseExceptionGroup) as cm: await runner() self.assertEqual(get_error_types(cm.exception), {MyBaseExc, TypeError}) async def test_taskgroup_22(self): async def foo1(): await asyncio.sleep(1) return 42 async def foo2(): await asyncio.sleep(2) return 11 async def runner(): async with taskgroups.TaskGroup() as g: g.create_task(foo1()) g.create_task(foo2()) r = asyncio.create_task(runner()) await asyncio.sleep(0.05) r.cancel() with self.assertRaises(asyncio.CancelledError): await r async def test_taskgroup_23(self): async def do_job(delay): await asyncio.sleep(delay) async with taskgroups.TaskGroup() as g: for count in range(10): await asyncio.sleep(0.1) g.create_task(do_job(0.3)) if count == 5: self.assertLess(len(g._tasks), 5) await asyncio.sleep(1.35) self.assertEqual(len(g._tasks), 0) async def test_taskgroup_24(self): async def root(g): await asyncio.sleep(0.1) g.create_task(coro1(0.1)) g.create_task(coro1(0.2)) async def coro1(delay): await asyncio.sleep(delay) async def runner(): async with taskgroups.TaskGroup() as g: g.create_task(root(g)) await runner() async def test_taskgroup_25(self): nhydras = 0 async def hydra(g): nonlocal nhydras nhydras += 1 await asyncio.sleep(0.01) g.create_task(hydra(g)) g.create_task(hydra(g)) async def hercules(): while nhydras < 10: await asyncio.sleep(0.015) 1 / 0 async def runner(): async with taskgroups.TaskGroup() as g: g.create_task(hydra(g)) g.create_task(hercules()) with self.assertRaises(ExceptionGroup) as cm: await runner() self.assertEqual(get_error_types(cm.exception), {ZeroDivisionError}) self.assertGreaterEqual(nhydras, 10) async def test_taskgroup_task_name(self): async def coro(): await asyncio.sleep(0) async with taskgroups.TaskGroup() as g: t = g.create_task(coro(), name="yolo") self.assertEqual(t.get_name(), "yolo") async def test_taskgroup_task_context(self): cvar = contextvars.ContextVar('cvar') async def coro(val): await asyncio.sleep(0) cvar.set(val) async with taskgroups.TaskGroup() as g: ctx = contextvars.copy_context() self.assertIsNone(ctx.get(cvar)) t1 = g.create_task(coro(1), context=ctx) await t1 self.assertEqual(1, ctx.get(cvar)) t2 = g.create_task(coro(2), context=ctx) await t2 self.assertEqual(2, ctx.get(cvar)) async def test_taskgroup_no_create_task_after_failure(self): async def coro1(): await asyncio.sleep(0.001) 1 / 0 async def coro2(g): try: await asyncio.sleep(1) except asyncio.CancelledError: with self.assertRaises(RuntimeError): g.create_task(coro1()) with self.assertRaises(ExceptionGroup) as cm: async with taskgroups.TaskGroup() as g: g.create_task(coro1()) g.create_task(coro2(g)) self.assertEqual(get_error_types(cm.exception), {ZeroDivisionError}) async def test_taskgroup_context_manager_exit_raises(self): # See https://github.com/python/cpython/issues/95289 class CustomException(Exception): pass async def raise_exc(): raise CustomException @contextlib.asynccontextmanager async def database(): try: yield finally: raise CustomException async def main(): task = asyncio.current_task() try: async with taskgroups.TaskGroup() as tg: async with database(): tg.create_task(raise_exc()) await asyncio.sleep(1) except* CustomException as err: self.assertEqual(task.cancelling(), 0) self.assertEqual(len(err.exceptions), 2) else: self.fail('CustomException not raised') await asyncio.create_task(main()) async def test_taskgroup_already_entered(self): tg = taskgroups.TaskGroup() async with tg: with self.assertRaisesRegex(RuntimeError, "has already been entered"): async with tg: pass async def test_taskgroup_double_enter(self): tg = taskgroups.TaskGroup() async with tg: pass with self.assertRaisesRegex(RuntimeError, "has already been entered"): async with tg: pass async def test_taskgroup_finished(self): async def create_task_after_tg_finish(): tg = taskgroups.TaskGroup() async with tg: pass coro = asyncio.sleep(0) with self.assertRaisesRegex(RuntimeError, "is finished"): tg.create_task(coro) # Make sure the coroutine was closed when submitted to the inactive tg # (if not closed, a RuntimeWarning should have been raised) with warnings.catch_warnings(record=True) as w: await create_task_after_tg_finish() self.assertEqual(len(w), 0) async def test_taskgroup_not_entered(self): tg = taskgroups.TaskGroup() coro = asyncio.sleep(0) with self.assertRaisesRegex(RuntimeError, "has not been entered"): tg.create_task(coro) async def test_taskgroup_without_parent_task(self): tg = taskgroups.TaskGroup() with self.assertRaisesRegex(RuntimeError, "parent task"): await await_without_task(tg.__aenter__()) coro = asyncio.sleep(0) with self.assertRaisesRegex(RuntimeError, "has not been entered"): tg.create_task(coro) def test_coro_closed_when_tg_closed(self): async def run_coro_after_tg_closes(): async with taskgroups.TaskGroup() as tg: pass coro = asyncio.sleep(0) with self.assertRaisesRegex(RuntimeError, "is finished"): tg.create_task(coro) loop = asyncio.get_event_loop() loop.run_until_complete(run_coro_after_tg_closes()) async def test_cancelling_level_preserved(self): async def raise_after(t, e): await asyncio.sleep(t) raise e() try: async with asyncio.TaskGroup() as tg: tg.create_task(raise_after(0.0, RuntimeError)) except* RuntimeError: pass self.assertEqual(asyncio.current_task().cancelling(), 0) async def test_nested_groups_both_cancelled(self): async def raise_after(t, e): await asyncio.sleep(t) raise e() try: async with asyncio.TaskGroup() as outer_tg: try: async with asyncio.TaskGroup() as inner_tg: inner_tg.create_task(raise_after(0, RuntimeError)) outer_tg.create_task(raise_after(0, ValueError)) except* RuntimeError: pass else: self.fail("RuntimeError not raised") self.assertEqual(asyncio.current_task().cancelling(), 1) except* ValueError: pass else: self.fail("ValueError not raised") self.assertEqual(asyncio.current_task().cancelling(), 0) async def test_error_and_cancel(self): event = asyncio.Event() async def raise_error(): event.set() await asyncio.sleep(0) raise RuntimeError() async def inner(): try: async with taskgroups.TaskGroup() as tg: tg.create_task(raise_error()) await asyncio.sleep(1) self.fail("Sleep in group should have been cancelled") except* RuntimeError: self.assertEqual(asyncio.current_task().cancelling(), 1) self.assertEqual(asyncio.current_task().cancelling(), 1) await asyncio.sleep(1) self.fail("Sleep after group should have been cancelled") async def outer(): t = asyncio.create_task(inner()) await event.wait() self.assertEqual(t.cancelling(), 0) t.cancel() self.assertEqual(t.cancelling(), 1) with self.assertRaises(asyncio.CancelledError): await t self.assertTrue(t.cancelled()) await outer() async def test_exception_refcycles_direct(self): """Test that TaskGroup doesn't keep a reference to the raised ExceptionGroup""" tg = asyncio.TaskGroup() exc = None class _Done(Exception): pass try: async with tg: raise _Done except ExceptionGroup as e: exc = e self.assertIsNotNone(exc) self.assertListEqual(gc.get_referrers(exc), no_other_refs()) async def test_exception_refcycles_errors(self): """Test that TaskGroup deletes self._errors, and __aexit__ args""" tg = asyncio.TaskGroup() exc = None class _Done(Exception): pass try: async with tg: raise _Done except* _Done as excs: exc = excs.exceptions[0] self.assertIsInstance(exc, _Done) self.assertListEqual(gc.get_referrers(exc), no_other_refs()) async def test_exception_refcycles_parent_task(self): """Test that TaskGroup deletes self._parent_task""" tg = asyncio.TaskGroup() exc = None class _Done(Exception): pass async def coro_fn(): async with tg: raise _Done try: async with asyncio.TaskGroup() as tg2: tg2.create_task(coro_fn()) except* _Done as excs: exc = excs.exceptions[0].exceptions[0] self.assertIsInstance(exc, _Done) self.assertListEqual(gc.get_referrers(exc), no_other_refs()) async def test_exception_refcycles_propagate_cancellation_error(self): """Test that TaskGroup deletes propagate_cancellation_error""" tg = asyncio.TaskGroup() exc = None try: async with asyncio.timeout(-1): async with tg: await asyncio.sleep(0) except TimeoutError as e: exc = e.__cause__ self.assertIsInstance(exc, asyncio.CancelledError) self.assertListEqual(gc.get_referrers(exc), no_other_refs()) async def test_exception_refcycles_base_error(self): """Test that TaskGroup deletes self._base_error""" class MyKeyboardInterrupt(KeyboardInterrupt): pass tg = asyncio.TaskGroup() exc = None try: async with tg: raise MyKeyboardInterrupt except MyKeyboardInterrupt as e: exc = e self.assertIsNotNone(exc) self.assertListEqual(gc.get_referrers(exc), no_other_refs()) if __name__ == "__main__": unittest.main()