From 9f1f4d803d9de1c7320d8dcad68456473276f6df Mon Sep 17 00:00:00 2001 From: George Li Date: Fri, 11 Apr 2025 19:21:57 +0000 Subject: [PATCH 1/2] asynq: fix asyncio mock Reviewers: dkang Test Plan: unit tests Differential Revision: Asana Tasks: Screenshot: CC: Deploy To: --- asynq/mock_.py | 33 +++++++++++++++++++++++++++++++-- asynq/tests/test_mock.py | 36 ++++++++++++++++++++++++++++++------ 2 files changed, 61 insertions(+), 8 deletions(-) diff --git a/asynq/mock_.py b/asynq/mock_.py index f561367..56cbb70 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,34 @@ 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..7735310 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())) # =================================================== From e2eb4db6bd85d6053f3d3023eac9026b78efbc41 Mon Sep 17 00:00:00 2001 From: George Li Date: Wed, 4 Jun 2025 05:38:58 +0000 Subject: [PATCH 2/2] black format changes --- asynq/mock_.py | 1 + asynq/tests/test_mock.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/asynq/mock_.py b/asynq/mock_.py index 56cbb70..86521ae 100644 --- a/asynq/mock_.py +++ b/asynq/mock_.py @@ -225,6 +225,7 @@ def __init__(self, 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): diff --git a/asynq/tests/test_mock.py b/asynq/tests/test_mock.py index 7735310..8b3820b 100644 --- a/asynq/tests/test_mock.py +++ b/asynq/tests/test_mock.py @@ -68,7 +68,7 @@ def class_method_non_async_caller(): async def class_method_asyncio_caller(): - return (await Cls.async_classmethod.asyncio()) + return await Cls.async_classmethod.asyncio() @asynq.asynq() @@ -83,7 +83,7 @@ def method_non_async_caller(): async def method_asyncio_caller(): - return (await Cls().async_method.asyncio()) + return await Cls().async_method.asyncio() class Counter(object):