bpo-37047: Refactor AsyncMock setup logic for autospeccing (GH-13574)
Handle late binding and attribute access in unittest.mock.AsyncMock
setup for autospeccing.
diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
index b14bf01..b91afd8 100644
--- a/Lib/unittest/mock.py
+++ b/Lib/unittest/mock.py
@@ -51,6 +51,13 @@
return False
+def _is_async_func(func):
+ if getattr(func, '__code__', None):
+ return asyncio.iscoroutinefunction(func)
+ else:
+ return False
+
+
def _is_instance_mock(obj):
# can't use isinstance on Mock objects because they override __class__
# The base class for all mocks is NonCallableMock
@@ -225,6 +232,34 @@
mock._mock_delegate = funcopy
+def _setup_async_mock(mock):
+ mock._is_coroutine = asyncio.coroutines._is_coroutine
+ mock.await_count = 0
+ mock.await_args = None
+ mock.await_args_list = _CallList()
+ mock.awaited = _AwaitEvent(mock)
+
+ # Mock is not configured yet so the attributes are set
+ # to a function and then the corresponding mock helper function
+ # is called when the helper is accessed similar to _setup_func.
+ def wrapper(attr, *args, **kwargs):
+ return getattr(mock.mock, attr)(*args, **kwargs)
+
+ for attribute in ('assert_awaited',
+ 'assert_awaited_once',
+ 'assert_awaited_with',
+ 'assert_awaited_once_with',
+ 'assert_any_await',
+ 'assert_has_awaits',
+ 'assert_not_awaited'):
+
+ # setattr(mock, attribute, wrapper) causes late binding
+ # hence attribute will always be the last value in the loop
+ # Use partial(wrapper, attribute) to ensure the attribute is bound
+ # correctly.
+ setattr(mock, attribute, partial(wrapper, attribute))
+
+
def _is_magic(name):
return '__%s__' % name[2:-2] == name
@@ -2151,7 +2186,7 @@
"""
self = _mock_self
if self.await_count != 0:
- msg = (f"Expected {self._mock_name or 'mock'} to have been awaited once."
+ msg = (f"Expected {self._mock_name or 'mock'} to not have been awaited."
f" Awaited {self.await_count} times.")
raise AssertionError(msg)
@@ -2457,10 +2492,7 @@
spec = type(spec)
is_type = isinstance(spec, type)
- if getattr(spec, '__code__', None):
- is_async_func = asyncio.iscoroutinefunction(spec)
- else:
- is_async_func = False
+ is_async_func = _is_async_func(spec)
_kwargs = {'spec': spec}
if spec_set:
_kwargs = {'spec_set': spec}
@@ -2498,26 +2530,11 @@
name=_name, **_kwargs)
if isinstance(spec, FunctionTypes):
- wrapped_mock = mock
# should only happen at the top level because we don't
# recurse for functions
mock = _set_signature(mock, spec)
if is_async_func:
- mock._is_coroutine = asyncio.coroutines._is_coroutine
- mock.await_count = 0
- mock.await_args = None
- mock.await_args_list = _CallList()
-
- for a in ('assert_awaited',
- 'assert_awaited_once',
- 'assert_awaited_with',
- 'assert_awaited_once_with',
- 'assert_any_await',
- 'assert_has_awaits',
- 'assert_not_awaited'):
- def f(*args, **kwargs):
- return getattr(wrapped_mock, a)(*args, **kwargs)
- setattr(mock, a, f)
+ _setup_async_mock(mock)
else:
_check_signature(spec, mock, is_type, instance)