You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
171 lines
6.2 KiB
171 lines
6.2 KiB
2 years ago
|
import asyncio
|
||
|
|
import inspect
|
||
|
|
|
||
|
|
from .case import TestCase
|
||
|
|
|
||
|
|
|
||
|
|
class IsolatedAsyncioTestCase(TestCase):
|
||
|
|
# Names intentionally have a long prefix
|
||
|
|
# to reduce a chance of clashing with user-defined attributes
|
||
|
|
# from inherited test case
|
||
|
|
#
|
||
|
|
# The class doesn't call loop.run_until_complete(self.setUp()) and family
|
||
|
|
# but uses a different approach:
|
||
|
|
# 1. create a long-running task that reads self.setUp()
|
||
|
|
# awaitable from queue along with a future
|
||
|
|
# 2. await the awaitable object passing in and set the result
|
||
|
|
# into the future object
|
||
|
|
# 3. Outer code puts the awaitable and the future object into a queue
|
||
|
|
# with waiting for the future
|
||
|
|
# The trick is necessary because every run_until_complete() call
|
||
|
|
# creates a new task with embedded ContextVar context.
|
||
|
|
# To share contextvars between setUp(), test and tearDown() we need to execute
|
||
|
|
# them inside the same task.
|
||
|
|
|
||
|
|
# Note: the test case modifies event loop policy if the policy was not instantiated
|
||
|
|
# yet.
|
||
|
|
# asyncio.get_event_loop_policy() creates a default policy on demand but never
|
||
|
|
# returns None
|
||
|
|
# I believe this is not an issue in user level tests but python itself for testing
|
||
|
|
# should reset a policy in every test module
|
||
|
|
# by calling asyncio.set_event_loop_policy(None) in tearDownModule()
|
||
|
|
|
||
|
|
def __init__(self, methodName='runTest'):
|
||
|
|
super().__init__(methodName)
|
||
|
|
self._asyncioTestLoop = None
|
||
|
|
self._asyncioCallsQueue = None
|
||
|
|
|
||
|
|
async def asyncSetUp(self):
|
||
|
|
pass
|
||
|
|
|
||
|
|
async def asyncTearDown(self):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def addAsyncCleanup(self, func, /, *args, **kwargs):
|
||
|
|
# A trivial trampoline to addCleanup()
|
||
|
|
# the function exists because it has a different semantics
|
||
|
|
# and signature:
|
||
|
|
# addCleanup() accepts regular functions
|
||
|
|
# but addAsyncCleanup() accepts coroutines
|
||
|
|
#
|
||
|
|
# We intentionally don't add inspect.iscoroutinefunction() check
|
||
|
|
# for func argument because there is no way
|
||
|
|
# to check for async function reliably:
|
||
|
|
# 1. It can be "async def func()" itself
|
||
|
|
# 2. Class can implement "async def __call__()" method
|
||
|
|
# 3. Regular "def func()" that returns awaitable object
|
||
|
|
self.addCleanup(*(func, *args), **kwargs)
|
||
|
|
|
||
|
|
def _callSetUp(self):
|
||
|
|
self.setUp()
|
||
|
|
self._callAsync(self.asyncSetUp)
|
||
|
|
|
||
|
|
def _callTestMethod(self, method):
|
||
|
|
self._callMaybeAsync(method)
|
||
|
|
|
||
|
|
def _callTearDown(self):
|
||
|
|
self._callAsync(self.asyncTearDown)
|
||
|
|
self.tearDown()
|
||
|
|
|
||
|
|
def _callCleanup(self, function, *args, **kwargs):
|
||
|
|
self._callMaybeAsync(function, *args, **kwargs)
|
||
|
|
|
||
|
|
def _callAsync(self, func, /, *args, **kwargs):
|
||
|
|
assert self._asyncioTestLoop is not None, 'asyncio test loop is not initialized'
|
||
|
|
ret = func(*args, **kwargs)
|
||
|
|
assert inspect.isawaitable(ret), f'{func!r} returned non-awaitable'
|
||
|
|
fut = self._asyncioTestLoop.create_future()
|
||
|
|
self._asyncioCallsQueue.put_nowait((fut, ret))
|
||
|
|
return self._asyncioTestLoop.run_until_complete(fut)
|
||
|
|
|
||
|
|
def _callMaybeAsync(self, func, /, *args, **kwargs):
|
||
|
|
assert self._asyncioTestLoop is not None, 'asyncio test loop is not initialized'
|
||
|
|
ret = func(*args, **kwargs)
|
||
|
|
if inspect.isawaitable(ret):
|
||
|
|
fut = self._asyncioTestLoop.create_future()
|
||
|
|
self._asyncioCallsQueue.put_nowait((fut, ret))
|
||
|
|
return self._asyncioTestLoop.run_until_complete(fut)
|
||
|
|
else:
|
||
|
|
return ret
|
||
|
|
|
||
|
|
async def _asyncioLoopRunner(self, fut):
|
||
|
|
self._asyncioCallsQueue = queue = asyncio.Queue()
|
||
|
|
fut.set_result(None)
|
||
|
|
while True:
|
||
|
|
query = await queue.get()
|
||
|
|
queue.task_done()
|
||
|
|
if query is None:
|
||
|
|
return
|
||
|
|
fut, awaitable = query
|
||
|
|
try:
|
||
|
|
ret = await awaitable
|
||
|
|
if not fut.cancelled():
|
||
|
|
fut.set_result(ret)
|
||
|
|
except (SystemExit, KeyboardInterrupt):
|
||
|
|
raise
|
||
|
|
except (BaseException, asyncio.CancelledError) as ex:
|
||
|
|
if not fut.cancelled():
|
||
|
|
fut.set_exception(ex)
|
||
|
|
|
||
|
|
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
|
||
|
|
fut = loop.create_future()
|
||
|
|
self._asyncioCallsTask = loop.create_task(self._asyncioLoopRunner(fut))
|
||
|
|
loop.run_until_complete(fut)
|
||
|
|
|
||
|
|
def _tearDownAsyncioLoop(self):
|
||
|
|
assert self._asyncioTestLoop is not None, 'asyncio test loop is not initialized'
|
||
|
|
loop = self._asyncioTestLoop
|
||
|
|
self._asyncioTestLoop = None
|
||
|
|
self._asyncioCallsQueue.put_nowait(None)
|
||
|
|
loop.run_until_complete(self._asyncioCallsQueue.join())
|
||
|
|
|
||
|
|
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:
|
||
|
|
# Prevent our executor environment from leaking to future tests.
|
||
|
|
loop.run_until_complete(loop.shutdown_default_executor())
|
||
|
|
asyncio.set_event_loop(None)
|
||
|
|
loop.close()
|
||
|
|
|
||
|
|
def run(self, result=None):
|
||
|
|
self._setupAsyncioLoop()
|
||
|
|
try:
|
||
|
|
return super().run(result)
|
||
|
|
finally:
|
||
|
|
self._tearDownAsyncioLoop()
|
||
|
|
|
||
|
|
def debug(self):
|
||
|
|
self._setupAsyncioLoop()
|
||
|
|
super().debug()
|
||
|
|
self._tearDownAsyncioLoop()
|
||
|
|
|
||
|
|
def __del__(self):
|
||
|
|
if self._asyncioTestLoop is not None:
|
||
|
|
self._tearDownAsyncioLoop()
|