diff --git a/asynq/mock_.py b/asynq/mock_.py index f561367..86521ae 100644 --- a/asynq/mock_.py +++ b/asynq/mock_.py @@ -148,9 +148,10 @@ def __enter__(self): mock_fn = super(_PatchAsync, self).__enter__() # so we can also mock non-functions for compatibility if callable(mock_fn): - async_fn = _AsyncWrapper(mock_fn) + async_fn = _AsynqWrapper(mock_fn) mock_fn.asynq = async_fn setattr(mock_fn, "async", async_fn) + mock_fn.asyncio = _AsyncioWrapper(mock_fn) return mock_fn def copy(self): @@ -185,7 +186,7 @@ def copy(self): return patcher -class _AsyncWrapper(object): +class _AsynqWrapper(object): """Wrapper for the .asynq attribute of patch'ed functions. Prevents people from setting and reading attributes on them. @@ -211,6 +212,35 @@ def __getattr__(self, attr): ) +class _AsyncioWrapper(object): + """Wrapper for the .asyncio attribute of patch'ed functions. + + Prevents people from setting and reading attributes on them. + + """ + + def __init__(self, mock_fn): + object.__setattr__(self, "_mock_fn", mock_fn) + + def __call__(self, *args, **kwargs): + async def async_wrapper(): + return self._mock_fn(*args, **kwargs) + + return async_wrapper() + + def __setattr__(self, attr, value): + raise TypeError( + "You cannot set attributes directly on a .asynq function. Set them on the " + "function itself instead." + ) + + def __getattr__(self, attr): + raise TypeError( + "You cannot read attributes directly on a .asynq function. Read them on the" + " function itself instead." + ) + + def _maybe_wrap_new(new): """If the mock replacement cannot have attributes set on it, wraps it in a function. diff --git a/asynq/tests/test_mock.py b/asynq/tests/test_mock.py index 3dd286f..8b3820b 100644 --- a/asynq/tests/test_mock.py +++ b/asynq/tests/test_mock.py @@ -13,6 +13,7 @@ # limitations under the License. import asynq +import asyncio from qcore.asserts import assert_eq, assert_is, assert_not_in, AssertRaises # =================================================== @@ -35,6 +36,10 @@ def non_async_caller(): return fn() +async def asyncio_caller(): + return await fn.asyncio() + + class Cls(object): @asynq.asynq() @classmethod @@ -62,6 +67,10 @@ def class_method_non_async_caller(): return Cls.async_classmethod() +async def class_method_asyncio_caller(): + return await Cls.async_classmethod.asyncio() + + @asynq.asynq() def method_async_caller(): obj = Cls() @@ -73,6 +82,10 @@ def method_non_async_caller(): return Cls().async_method() +async def method_asyncio_caller(): + return await Cls().async_method.asyncio() + + class Counter(object): linked_class = Cls @@ -97,20 +110,28 @@ class SecondClass(object): class MockChecker(object): @classmethod def check(cls, mock_fn, mock_classmethod, mock_method): - cls._check_mock(mock_fn, async_caller, non_async_caller) + cls._check_mock(mock_fn, async_caller, non_async_caller, asyncio_caller) + cls._check_mock( + mock_classmethod, + class_method_async_caller, + class_method_non_async_caller, + class_method_asyncio_caller, + ) cls._check_mock( - mock_classmethod, class_method_async_caller, class_method_non_async_caller + mock_method, + method_async_caller, + method_non_async_caller, + method_asyncio_caller, ) - cls._check_mock(mock_method, method_async_caller, method_non_async_caller) @classmethod - def _check_mock(cls, mock_fn, async_caller, non_async_caller): + def _check_mock(cls, mock_fn, async_caller, non_async_caller, asyncio_caller): raise NotImplementedError class MockCheckerWithAssignment(MockChecker): @classmethod - def _check_mock(cls, mock_fn, async_caller, non_async_caller): + def _check_mock(cls, mock_fn, async_caller, non_async_caller, asyncio_caller): mock_fn.return_value = 42 assert_eq(0, mock_fn.call_count) assert_eq(42, non_async_caller()) @@ -123,12 +144,15 @@ def _check_mock(cls, mock_fn, async_caller, non_async_caller): assert_eq(43, non_async_caller()) assert_eq((43, 43), async_caller()) + assert_eq(43, asyncio.run(asyncio_caller())) + class MockCheckerWithNew(MockChecker): @classmethod - def _check_mock(cls, mock_fn, async_caller, non_async_caller): + def _check_mock(cls, mock_fn, async_caller, non_async_caller, asyncio_caller): assert_eq(42, non_async_caller()) assert_eq((42, 42), async_caller()) + assert_eq(42, asyncio.run(asyncio_caller())) # ===================================================