mock/backports.py

import sys


if sys.version_info[:2] < (3, 8):

    import asyncio, functools
    from asyncio.coroutines import _is_coroutine
    from inspect import ismethod, isfunction, CO_COROUTINE
    from unittest import TestCase

    def _unwrap_partial(func):
        while isinstance(func, functools.partial):
            func = func.func
        return func

    def _has_code_flag(f, flag):
        """Return true if ``f`` is a function (or a method or functools.partial
        wrapper wrapping a function) whose code object has the given ``flag``
        set in its flags."""
        while ismethod(f):
            f = f.__func__
        f = _unwrap_partial(f)
        if not isfunction(f):
            return False
        return bool(f.__code__.co_flags & flag)

    def iscoroutinefunction(obj):
        """Return true if the object is a coroutine function.

        Coroutine functions are defined with "async def" syntax.
        """
        return (
            _has_code_flag(obj, CO_COROUTINE) or
            getattr(obj, '_is_coroutine', None) is _is_coroutine
        )


    class IsolatedAsyncioTestCase(TestCase):

        def __init__(self, methodName='runTest'):
            super().__init__(methodName)
            self._asyncioTestLoop = None
            self._asyncioCallsQueue = None

        async def _asyncioLoopRunner(self, fut):
            self._asyncioCallsQueue = queue = asyncio.Queue()
            fut.set_result(None)
            while True:
                query = await queue.get()
                queue.task_done()
                assert query is None

        def _setupAsyncioLoop(self):
            assert self._asyncioTestLoop is None
            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
            loop = self._asyncioTestLoop
            self._asyncioTestLoop = None
            self._asyncioCallsQueue.put_nowait(None)
            loop.run_until_complete(self._asyncioCallsQueue.join())

            try:
                # shutdown asyncgens
                loop.run_until_complete(loop.shutdown_asyncgens())
            finally:
                asyncio.set_event_loop(None)
                loop.close()

        def run(self, result=None):
            self._setupAsyncioLoop()
            try:
                return super().run(result)
            finally:
                self._tearDownAsyncioLoop()


else:

    from asyncio import iscoroutinefunction
    from unittest import IsolatedAsyncioTestCase