diff --git a/core/tests/mock.py b/core/tests/mock.py index f5dcb7d484..810d2c8a38 100644 --- a/core/tests/mock.py +++ b/core/tests/mock.py @@ -24,6 +24,19 @@ class Mock: return self.return_value +class MockAsync(Mock): + def __call__(self, *args, **kwargs) -> Any: + self.calls.append((args, kwargs)) + + if self.raises is not None: + raise self.raises + + async def _async_return() -> Any: + return self.return_value + + return _async_return() + + class patch: MOCK_OBJECT = object() NO_VALUE = object()