bpo-47062: Implement asyncio.Runner context manager (GH-31799)

Co-authored-by: Zachary Ware <zach@python.org>

1 files changed, 19 insertions, 51 deletions

diff --git a/Lib/unittest/async_case.py b/Lib/unittest/async_case.py
index 25adc3d..85b938f 100644
--- a/Lib/unittest/async_case.py
+++ b/Lib/unittest/async_case.py
@@ -34,7 +34,7 @@ class IsolatedAsyncioTestCase(TestCase):
 
     def __init__(self, methodName='runTest'):
         super().__init__(methodName)
-        self._asyncioTestLoop = None
+        self._asyncioRunner = None
         self._asyncioTestContext = contextvars.copy_context()
 
     async def asyncSetUp(self):
@@ -75,76 +75,44 @@ class IsolatedAsyncioTestCase(TestCase):
         self._callMaybeAsync(function, *args, **kwargs)
 
     def _callAsync(self, func, /, *args, **kwargs):
-        assert self._asyncioTestLoop is not None, 'asyncio test loop is not initialized'
+        assert self._asyncioRunner is not None, 'asyncio runner is not initialized'
         assert inspect.iscoroutinefunction(func), f'{func!r} is not an async function'
-        task = self._asyncioTestLoop.create_task(
+        return self._asyncioRunner.run(
             func(*args, **kwargs),
-            context=self._asyncioTestContext,
+            context=self._asyncioTestContext
         )
-        return self._asyncioTestLoop.run_until_complete(task)
 
     def _callMaybeAsync(self, func, /, *args, **kwargs):
-        assert self._asyncioTestLoop is not None, 'asyncio test loop is not initialized'
+        assert self._asyncioRunner is not None, 'asyncio runner is not initialized'
         if inspect.iscoroutinefunction(func):
-            task = self._asyncioTestLoop.create_task(
+            return self._asyncioRunner.run(
                 func(*args, **kwargs),
                 context=self._asyncioTestContext,
             )
-            return self._asyncioTestLoop.run_until_complete(task)
         else:
             return self._asyncioTestContext.run(func, *args, **kwargs)
 
-    def _setupAsyncioLoop(self):
-        assert self._asyncioTestLoop is None, 'asyncio test loop already initialized'
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-        loop.set_debug(True)
-        self._asyncioTestLoop = loop
+    def _setupAsyncioRunner(self):
+        assert self._asyncioRunner is None, 'asyncio runner is already initialized'
+        runner = asyncio.Runner(debug=True)
+        self._asyncioRunner = runner
 
-    def _tearDownAsyncioLoop(self):
-        assert self._asyncioTestLoop is not None, 'asyncio test loop is not initialized'
-        loop = self._asyncioTestLoop
-        self._asyncioTestLoop = None
-
-        try:
-            # cancel all tasks
-            to_cancel = asyncio.all_tasks(loop)
-            if not to_cancel:
-                return
-
-            for task in to_cancel:
-                task.cancel()
-
-            loop.run_until_complete(
-                asyncio.gather(*to_cancel, return_exceptions=True))
-
-            for task in to_cancel:
-                if task.cancelled():
-                    continue
-                if task.exception() is not None:
-                    loop.call_exception_handler({
-                        'message': 'unhandled exception during test shutdown',
-                        'exception': task.exception(),
-                        'task': task,
-                    })
-            # shutdown asyncgens
-            loop.run_until_complete(loop.shutdown_asyncgens())
-        finally:
-            asyncio.set_event_loop(None)
-            loop.close()
+    def _tearDownAsyncioRunner(self):
+        runner = self._asyncioRunner
+        runner.close()
 
     def run(self, result=None):
-        self._setupAsyncioLoop()
+        self._setupAsyncioRunner()
         try:
             return super().run(result)
         finally:
-            self._tearDownAsyncioLoop()
+            self._tearDownAsyncioRunner()
 
     def debug(self):
-        self._setupAsyncioLoop()
+        self._setupAsyncioRunner()
         super().debug()
-        self._tearDownAsyncioLoop()
+        self._tearDownAsyncioRunner()
 
     def __del__(self):
-        if self._asyncioTestLoop is not None:
-            self._tearDownAsyncioLoop()
+        if self._asyncioRunner is not None:
+            self._tearDownAsyncioRunner()