asynq: fix asyncio mock

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.
 

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()))
 
 
 # ===================================================