anyio.py

#  Copyright 2019  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
"""

from functools import wraps
from unittest import mock

import anyio


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

	def decorator(testfunc):
		async def test_async_wrapper(args):
			async with anyio.fail_after(timeout):
				return await testfunc(*args)

		@wraps(testfunc)
		def test_wrapper(*args):
			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, delay=0, **kwargs):
		mock._safe_super(AsyncMock, self).__init__(*args, **kwargs)
		self.delay = delay

	async def __call__(_mock_self, *args, **kwargs):
		_mock_self._mock_check_sig(*args, **kwargs)
		await anyio.sleep(_mock_self.delay)
		return _mock_self._mock_call(*args, **kwargs)