jpayne@69: import asyncio jpayne@69: import inspect jpayne@69: jpayne@69: from .case import TestCase jpayne@69: jpayne@69: jpayne@69: jpayne@69: class IsolatedAsyncioTestCase(TestCase): jpayne@69: # Names intentionally have a long prefix jpayne@69: # to reduce a chance of clashing with user-defined attributes jpayne@69: # from inherited test case jpayne@69: # jpayne@69: # The class doesn't call loop.run_until_complete(self.setUp()) and family jpayne@69: # but uses a different approach: jpayne@69: # 1. create a long-running task that reads self.setUp() jpayne@69: # awaitable from queue along with a future jpayne@69: # 2. await the awaitable object passing in and set the result jpayne@69: # into the future object jpayne@69: # 3. Outer code puts the awaitable and the future object into a queue jpayne@69: # with waiting for the future jpayne@69: # The trick is necessary because every run_until_complete() call jpayne@69: # creates a new task with embedded ContextVar context. jpayne@69: # To share contextvars between setUp(), test and tearDown() we need to execute jpayne@69: # them inside the same task. jpayne@69: jpayne@69: # Note: the test case modifies event loop policy if the policy was not instantiated jpayne@69: # yet. jpayne@69: # asyncio.get_event_loop_policy() creates a default policy on demand but never jpayne@69: # returns None jpayne@69: # I believe this is not an issue in user level tests but python itself for testing jpayne@69: # should reset a policy in every test module jpayne@69: # by calling asyncio.set_event_loop_policy(None) in tearDownModule() jpayne@69: jpayne@69: def __init__(self, methodName='runTest'): jpayne@69: super().__init__(methodName) jpayne@69: self._asyncioTestLoop = None jpayne@69: self._asyncioCallsQueue = None jpayne@69: jpayne@69: async def asyncSetUp(self): jpayne@69: pass jpayne@69: jpayne@69: async def asyncTearDown(self): jpayne@69: pass jpayne@69: jpayne@69: def addAsyncCleanup(self, func, /, *args, **kwargs): jpayne@69: # A trivial trampoline to addCleanup() jpayne@69: # the function exists because it has a different semantics jpayne@69: # and signature: jpayne@69: # addCleanup() accepts regular functions jpayne@69: # but addAsyncCleanup() accepts coroutines jpayne@69: # jpayne@69: # We intentionally don't add inspect.iscoroutinefunction() check jpayne@69: # for func argument because there is no way jpayne@69: # to check for async function reliably: jpayne@69: # 1. It can be "async def func()" iself jpayne@69: # 2. Class can implement "async def __call__()" method jpayne@69: # 3. Regular "def func()" that returns awaitable object jpayne@69: self.addCleanup(*(func, *args), **kwargs) jpayne@69: jpayne@69: def _callSetUp(self): jpayne@69: self.setUp() jpayne@69: self._callAsync(self.asyncSetUp) jpayne@69: jpayne@69: def _callTestMethod(self, method): jpayne@69: self._callMaybeAsync(method) jpayne@69: jpayne@69: def _callTearDown(self): jpayne@69: self._callAsync(self.asyncTearDown) jpayne@69: self.tearDown() jpayne@69: jpayne@69: def _callCleanup(self, function, *args, **kwargs): jpayne@69: self._callMaybeAsync(function, *args, **kwargs) jpayne@69: jpayne@69: def _callAsync(self, func, /, *args, **kwargs): jpayne@69: assert self._asyncioTestLoop is not None jpayne@69: ret = func(*args, **kwargs) jpayne@69: assert inspect.isawaitable(ret) jpayne@69: fut = self._asyncioTestLoop.create_future() jpayne@69: self._asyncioCallsQueue.put_nowait((fut, ret)) jpayne@69: return self._asyncioTestLoop.run_until_complete(fut) jpayne@69: jpayne@69: def _callMaybeAsync(self, func, /, *args, **kwargs): jpayne@69: assert self._asyncioTestLoop is not None jpayne@69: ret = func(*args, **kwargs) jpayne@69: if inspect.isawaitable(ret): jpayne@69: fut = self._asyncioTestLoop.create_future() jpayne@69: self._asyncioCallsQueue.put_nowait((fut, ret)) jpayne@69: return self._asyncioTestLoop.run_until_complete(fut) jpayne@69: else: jpayne@69: return ret jpayne@69: jpayne@69: async def _asyncioLoopRunner(self, fut): jpayne@69: self._asyncioCallsQueue = queue = asyncio.Queue() jpayne@69: fut.set_result(None) jpayne@69: while True: jpayne@69: query = await queue.get() jpayne@69: queue.task_done() jpayne@69: if query is None: jpayne@69: return jpayne@69: fut, awaitable = query jpayne@69: try: jpayne@69: ret = await awaitable jpayne@69: if not fut.cancelled(): jpayne@69: fut.set_result(ret) jpayne@69: except asyncio.CancelledError: jpayne@69: raise jpayne@69: except Exception as ex: jpayne@69: if not fut.cancelled(): jpayne@69: fut.set_exception(ex) jpayne@69: jpayne@69: def _setupAsyncioLoop(self): jpayne@69: assert self._asyncioTestLoop is None jpayne@69: loop = asyncio.new_event_loop() jpayne@69: asyncio.set_event_loop(loop) jpayne@69: loop.set_debug(True) jpayne@69: self._asyncioTestLoop = loop jpayne@69: fut = loop.create_future() jpayne@69: self._asyncioCallsTask = loop.create_task(self._asyncioLoopRunner(fut)) jpayne@69: loop.run_until_complete(fut) jpayne@69: jpayne@69: def _tearDownAsyncioLoop(self): jpayne@69: assert self._asyncioTestLoop is not None jpayne@69: loop = self._asyncioTestLoop jpayne@69: self._asyncioTestLoop = None jpayne@69: self._asyncioCallsQueue.put_nowait(None) jpayne@69: loop.run_until_complete(self._asyncioCallsQueue.join()) jpayne@69: jpayne@69: try: jpayne@69: # cancel all tasks jpayne@69: to_cancel = asyncio.all_tasks(loop) jpayne@69: if not to_cancel: jpayne@69: return jpayne@69: jpayne@69: for task in to_cancel: jpayne@69: task.cancel() jpayne@69: jpayne@69: loop.run_until_complete( jpayne@69: asyncio.gather(*to_cancel, loop=loop, return_exceptions=True)) jpayne@69: jpayne@69: for task in to_cancel: jpayne@69: if task.cancelled(): jpayne@69: continue jpayne@69: if task.exception() is not None: jpayne@69: loop.call_exception_handler({ jpayne@69: 'message': 'unhandled exception during test shutdown', jpayne@69: 'exception': task.exception(), jpayne@69: 'task': task, jpayne@69: }) jpayne@69: # shutdown asyncgens jpayne@69: loop.run_until_complete(loop.shutdown_asyncgens()) jpayne@69: finally: jpayne@69: asyncio.set_event_loop(None) jpayne@69: loop.close() jpayne@69: jpayne@69: def run(self, result=None): jpayne@69: self._setupAsyncioLoop() jpayne@69: try: jpayne@69: return super().run(result) jpayne@69: finally: jpayne@69: self._tearDownAsyncioLoop()