_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 Tuple
from unittest import mock

import anyio

py_version = sys.version_info[:2]

AnyFunc = Callable[..., Any]
Decorator = Callable[[AnyFunc], AnyFunc]


def with_anyio(timeout: int = 10) -> Decorator:
	"""
	Create a wrapping decorator to run asynchronous test functions
	"""

	def decorator(testfunc: AnyFunc) -> AnyFunc:
		async def test_async_wrapper(args: Tuple[Any]) -> Any:
			async with anyio.fail_after(timeout):
				return await testfunc(*args)

		@wraps(testfunc)
		def test_wrapper(*args: Any) -> Any:
			return anyio.run(test_async_wrapper, args)

		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)