_anyio.py

#  Copyright 2019-2021  Dom Sekotill <dom.sekotill@kodo.org.uk>
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""
Anyio helpers for unit tests
"""

import sys
from functools import wraps
from typing import Any
from typing import Callable
from typing import Coroutine
from typing import Literal
from typing import Tuple
from typing import Union
from unittest import TestCase
from unittest import mock
from warnings import warn

import anyio

try:
	import trio as _  # noqa
	USE_TRIO = True
except ImportError:
	USE_TRIO = False

Backend = Union[Literal['asyncio'], Literal['trio']]

py_version = sys.version_info[:2]

AsyncTestFunc = Callable[..., Coroutine[Any, Any, None]]
TestFunc = Callable[..., None]


def with_anyio(*backends: Backend, timeout: int = 10) -> Callable[[AsyncTestFunc], TestFunc]:
	"""
	Create a wrapping decorator to run asynchronous test functions
	"""
	if not backends:
		backends = ('asyncio',)

	def decorator(testfunc: AsyncTestFunc) -> TestFunc:
		async def test_async_wrapper(tc: TestCase, args: Tuple[mock.Mock]) -> None:
			with anyio.fail_after(timeout):
				await testfunc(tc, *args)

		@wraps(testfunc)
		def test_wrapper(tc: TestCase, *args: mock.Mock) -> None:
			for backend in backends:
				if backend == 'trio' and not USE_TRIO:
					warn(
						f"not running {testfunc.__name__} with trio; package is missing",
					)
					continue
				with tc.subTest(f"backend: {backend}"):
					anyio.run(test_async_wrapper, tc, args, backend=backend)

		return test_wrapper

	return decorator


class AsyncMock(mock.Mock):
	"""
	A Mock class that acts as a coroutine when called
	"""

	def __init__(self, *args: Any, delay: float = 0.0, **kwargs: Any):
		mock._safe_super(AsyncMock, self).__init__(*args, **kwargs)  # type: ignore
		self.delay = delay

	async def __call__(_mock_self, *args: Any, **kwargs: Any) -> Any:
		_mock_self._mock_check_sig(*args, **kwargs)
		if py_version >= (3, 8):
			_mock_self._increment_mock_call(*args, **kwargs)
		await anyio.sleep(_mock_self.delay)
		return _mock_self._mock_call(*args, **kwargs)