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