Adds AsyncMock docs and tests for iterators and context managers.

python · lisroach · May 20, 2019 · May 14, 2018 · Jul 22, 2018 · Aug 16, 2018
commit ae13db16c2469ef7471e0087911f6bd796d3d965
diff --git a/Doc/library/unittest.mock.rst b/Doc/library/unittest.mock.rst
@@ -851,6 +851,165 @@ object::
     >>> p.assert_called_once_with()
 
 
+.. class:: AsyncMock(spec=None, side_effect=None, return_value=DEFAULT, wraps=None, name=None, spec_set=None, unsafe=False, **kwargs)
+
+  An asynchronous version of :class:`Mock`. The :class:`AsyncMock` object will
+  behave so the object is recognized as an async function, and the result of a
+  call is an awaitable.
+
+    >>> mock = AsyncMock()
+    >>> asyncio.iscoroutinefunction(mock)
+    True
+    >>> inspect.isawaitable(mock())
+    True
+
+  The result of ``mock()`` is an async function which will have the outcome
+  of ``side_effect`` or ``return_value``:
+
+  - if ``side_effect`` is a function, the async function will return the
+    result of that function,
+  - if ``side_effect`` is an exception, the async function will raise the
+    exception,
+  - if ``side_effect`` is an iterable, the async function will return the
+    next value of the iterable, however, if the sequence of result is
+    exhausted, ``StopIteration`` is raised immediately,
+  - if ``side_effect`` is not defined, the async function will return the
+    value defined by ``return_value``, hence, by default, the async function
+    returns a new :class:`AsyncMock` object.
+
+
+  Setting the *spec* of a :class:`Mock` or :class:`MagicMock` to an async function
+  will result in a coroutine object being returned after calling.
+
+    >>> async def async_func(): pass
+    ...
+    >>> mock = MagicMock(async_func)
+    >>> mock
+    <MagicMock spec='function' id='4568403696'>
+    >>> mock()
+    <coroutine object AsyncMockMixin._mock_call at 0x1104cb440>
+
+  .. method:: assert_awaited()
+
+      Assert that the mock was awaited at least once.
+
+          >>> mock = AsyncMock()
+          >>> async def main():
+          ...     await mock()
+          ...
+          >>> asyncio.run(main())
+          >>> mock.assert_awaited()
+          >>> mock_2 = AsyncMock()
+          >>> mock_2.assert_awaited()
+          Traceback (most recent call last):
+          ...
+          AssertionError: Expected mock to have been awaited.
+
+  .. method:: assert_awaited_once()
+
+      Assert that the mock was awaited exactly once.
+
+        >>> mock = AsyncMock()
+        >>> async def main():
+        ...     await mock()
+        ...
+        >>> asyncio.run(main())
+        >>> mock.assert_awaited_once()
+        >>> asyncio.run(main())
+        >>> mock.method.assert_awaited_once()
+        Traceback (most recent call last):
+        ...
+        AssertionError: Expected mock to have been awaited once. Awaited 2 times.
+
+  .. method:: assert_awaited_with(*args, **kwargs)
+
+      Assert that the last await was with the specified arguments.
+
+        >>> mock = AsyncMock()
+        >>> async def main(*args, **kwargs):
+        ...     await mock(*args, **kwargs)
+        ...
+        >>> asyncio.run(main('foo', bar='bar'))
+        >>> mock.assert_awaited_with('foo', bar='bar')
+        >>> mock.assert_awaited_with('other')
+        Traceback (most recent call last):
+        ...
+        AssertionError: expected call not found.
+        Expected: mock('other')
+        Actual: mock('foo', bar='bar')
+
+  .. method:: assert_awaited_once_with(*args, **kwargs)
+
+      Assert that the mock was awaited exactly once and with the specified
+      arguments.
+
+        >>> mock = AsyncMock()
+        >>> async def main(*args, **kwargs):
+        ...     await mock(*args, **kwargs)
+        ...
+        >>> asyncio.run(main('foo', bar='bar'))
+        >>> mock.assert_awaited_once_with('foo', bar='bar')
+        >>> asyncio.run(main('foo', bar='bar'))
+        >>> mock.assert_awaited_once_with('foo', bar='bar')
+        Traceback (most recent call last):
+        ...
+        AssertionError: Expected mock to have been awaited once. Awaited 2 times.
+
+  .. method:: assert_any_await(*args, **kwargs)
+
+      Assert the mock has ever been awaited with the specified arguments.
+
+        >>> mock = AsyncMock()
+        >>> async def main(*args, **kwargs):
+        ...     await mock(*args, **kwargs)
+        ...
+        >>> asyncio.run(main('foo', bar='bar'))
+        >>> asyncio.run(main('hello'))
+        >>> mock.assert_any_await('foo', bar='bar')
+        >>> mock.assert_any_await('other')
+        Traceback (most recent call last):
+        ...
+        AssertionError: mock('other') await not found
+
+  .. method:: assert_has_awaits(calls, any_order=False)
+
+      Assert the mock has been awaited with the specified calls.
+      The :attr:`await_args_list` list is checked for the awaits.
+
+      If `any_order` is False (the default) then the awaits must be
+      sequential. There can be extra calls before or after the
+      specified awaits.
+
+      If `any_order` is True then the awaits can be in any order, but
+      they must all appear in :attr:`await_args_list`.
+
+        >>> mock = AsyncMock()
+        >>> async def main(*args, **kwargs):
+        ...     await mock(*args, **kwargs)
+        ...
+        >>> calls = [call("foo"), call("bar")]
+        >>> mock.assert_has_calls(calls)
+        Traceback (most recent call last):
+        ...
+        AssertionError: Calls not found.
+        Expected: [call('foo'), call('bar')]
+        >>> asyncio.run(main('foo'))
+        >>> asyncio.run(main('bar'))
+        >>> mock.assert_has_calls(calls)
+
+  .. method:: assert_not_awaited()
+
+    Assert that the mock was never awaited.
+
+        >>> mock = AsyncMock()
+        >>> mock.assert_not_awaited()
+
+  .. method:: reset_mock(*args, **kwargs)
+
+    See :func:`Mock.reset_mock`. Also sets :attr:`await_count` to 0,
+    :attr:`await_args` to None, and clears the :attr:`await_args_list`.
+
+
 Calling
 ~~~~~~~
 

diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
@@ -1818,6 +1818,7 @@ def method(self, *args, **kw):
     '__float__': 1.0,
     '__bool__': True,
     '__index__': 1,
+    '__aexit__': False,
 }
 
 
@@ -1854,8 +1855,8 @@ def _get_async_iter(self):
     def __aiter__():
         ret_val = self.__aiter__._mock_return_value
         if ret_val is DEFAULT:
-            return AsyncIterator([])
-        return AsyncIterator(ret_val)
+            return _AsyncIterator(iter([]))
+        return _AsyncIterator(iter(ret_val))
     return __aiter__
 
 _side_effect_methods = {
@@ -1927,8 +1928,33 @@ def mock_add_spec(self, spec, spec_set=False):
         self._mock_set_magics()
 
 
+class AsyncMagicMixin:
+    def __init__(self, *args, **kw):
+        self._mock_set_async_magics()  # make magic work for kwargs in init
+        _safe_super(AsyncMagicMixin, self).__init__(*args, **kw)
+        self._mock_set_async_magics()  # fix magic broken by upper level init
+
+    def _mock_set_async_magics(self):
+        these_magics = _async_magics
+
+        if getattr(self, "_mock_methods", None) is not None:
+            these_magics = _async_magics.intersection(self._mock_methods)
+            remove_magics = _async_magics - these_magics
+
+            for entry in remove_magics:
+                if entry in type(self).__dict__:
+                    # remove unneeded magic methods
+                    delattr(self, entry)
+
+        # don't overwrite existing attributes if called a second time
+        these_magics = these_magics - set(type(self).__dict__)
 
-class MagicMock(MagicMixin, Mock):
+        _type = type(self)
+        for entry in these_magics:
+            setattr(_type, entry, MagicProxy(entry, self))
+
+
+class MagicMock(MagicMixin, AsyncMagicMixin, Mock):
     """
     MagicMock is a subclass of Mock with default implementations
     of most of the magic methods. You can use MagicMock without having to
@@ -1968,32 +1994,6 @@ def __get__(self, obj, _type=None):
         return self.create_mock()
 
 
-class AsyncMagicMixin:
-    def __init__(self, *args, **kw):
-        self._mock_set_async_magics()  # make magic work for kwargs in init
-        _safe_super(AsyncMagicMixin, self).__init__(*args, **kw)
-        self._mock_set_async_magics()  # fix magic broken by upper level init
-
-    def _mock_set_async_magics(self):
-        these_magics = _async_magics
-
-        if getattr(self, "_mock_methods", None) is not None:
-            these_magics = _async_magics.intersection(self._mock_methods)
-            remove_magics = _async_magics - these_magics
-
-            for entry in remove_magics:
-                if entry in type(self).__dict__:
-                    # remove unneeded magic methods
-                    delattr(self, entry)
-
-        # don't overwrite existing attributes if called a second time
-        these_magics = these_magics - set(type(self).__dict__)
-
-        _type = type(self)
-        for entry in these_magics:
-            setattr(_type, entry, MagicProxy(entry, self))
-
-
 class AsyncMockMixin(Base):
     awaited = _delegating_property('awaited')
     await_count = _delegating_property('await_count')
@@ -2052,13 +2052,13 @@ def assert_awaited(_mock_self):
             msg = f"Expected {self._mock_name or 'mock'} to have been awaited."
             raise AssertionError(msg)
 
-    def assert_awaited_once(_mock_self, *args, **kwargs):
+    def assert_awaited_once(_mock_self):
         """
         Assert that the mock was awaited exactly once.
         """
         self = _mock_self
         if not self.await_count == 1:
-            msg = (f"Expected {self._mock_name or 'mock'} to have been awaited once.",
+            msg = (f"Expected {self._mock_name or 'mock'} to have been awaited once."
                    f" Awaited {self.await_count} times.")
             raise AssertionError(msg)
 
@@ -2088,7 +2088,7 @@ def assert_awaited_once_with(_mock_self, *args, **kwargs):
         """
         self = _mock_self
         if not self.await_count == 1:
-            msg = (f"Expected {self._mock_name or 'mock'} to have been awaited once.",
+            msg = (f"Expected {self._mock_name or 'mock'} to have been awaited once."
                    f" Awaited {self.await_count} times.")
             raise AssertionError(msg)
         return self.assert_awaited_with(*args, **kwargs)
@@ -2150,7 +2150,7 @@ def assert_not_awaited(_mock_self):
         """
         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 have been awaited once."
                    f" Awaited {self.await_count} times.")
             raise AssertionError(msg)
 
@@ -2167,10 +2167,10 @@ def reset_mock(self, *args, **kwargs):
 class AsyncMock(AsyncMockMixin, AsyncMagicMixin, Mock):
     """
     Enhance :class:`Mock` with features allowing to mock
-    a async function.
+    an async function.
 
     The :class:`AsyncMock` object will behave so the object is
-    recognized as async function, and the result of a call as a async:
+    recognized as an async function, and the result of a call is an awaitable:
 
     >>> mock = AsyncMock()
     >>> asyncio.iscoroutinefunction(mock)
@@ -2193,8 +2193,8 @@ class AsyncMock(AsyncMockMixin, AsyncMagicMixin, Mock):
       value defined by ``return_value``, hence, by default, the async function
       returns a new :class:`AsyncMock` object.
 
-    If the outcome of ``side_effect`` or ``return_value`` is an async function, the
-    mock async function obtained when the mock object is called will be this
+    If the outcome of ``side_effect`` or ``return_value`` is an async function,
+    the mock async function obtained when the mock object is called will be this
     async function itself (and not an async function returning an async
     function).
 
@@ -2757,20 +2757,19 @@ async def _raise(exception):
     raise exception
 
 
-class AsyncIterator:
+class _AsyncIterator:
     """
     Wraps an iterator in an asynchronous iterator.
     """
     def __init__(self, iterator):
         self.iterator = iterator
         code_mock = NonCallableMock(spec_set=CodeType)
-        code_mock.co_flags = inspect.CO_ITERATOR_COROUTINE
+        code_mock.co_flags = inspect.CO_ITERABLE_COROUTINE
         self.__dict__['__code__'] = code_mock
 
     def __aiter__(self):
         return self
 
-    # cannot use async before 3.7
     async def __anext__(self):
         try:
             return next(self.iterator)