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

Co-authored-by: Zachary Ware <zach@python.org>
This commit is contained in:
Andrew Svetlov 2022-03-24 21:51:16 +02:00 committed by GitHub
parent 2f49b97cc5
commit 4119d2d7c9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 381 additions and 106 deletions

View file

@ -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()