Alternative to mock_calls for awaits in AsyncMock

Ideas
1

I’d like to have an alternative to mock_calls for awaited coroutines to AsyncMock.

Mock.mock_calls contain all calls to the mock and its children. When used on AsyncMock, the mock_calls record is added when coroutine is created, regardless of whether it was awaited. Awaits can be checked using assert_await* methods or await_* attributes, but those only record the mock itself and there is no option to check awaits on the AsyncMock and its children at once.

The best case scenario, it should be possible to switch from sync version

mock = Mock()
mock(sentinel.foo)
assert mock.mock_calls == [call(sentinel.foo)]

to async version

mock = AsyncMock()
await mock(sentinel.foo)
assert mock.mock_awaits == [call(sentinel.foo)]

Since I commonly use mock_calls checks in tests and there is no alternative for awaits, I tend to use it also for async code and omit the checks for awaits. I’m afraid that I may miss some awaits this way.

2 Likes
2

Here’s a utility function as workaround:

from typing import List, Tuple
from unittest.mock import AsyncMock, Mock, _Call

def mock_awaits(mock: AsyncMock) -> List[_Call]:
    """Return a list of awaited call from mock and all its children.

    Awaits are in an order of their coroutine calls.
    """
    # A list of awaited calls.
    awaits = []
    # A queue of mock objects to process.
    queue: List[Tuple[str, Mock]] = []
    queue.append(('', mock))
    # Collect awaits from the structure of mocks.
    # Based on `CallableMixin._increment_mock_call`.
    while queue:
        parent_name, a_mock = queue.pop(0)

        # Collect children to process.
        children = []
        if isinstance(a_mock._mock_return_value, Mock):
            children.append(a_mock._mock_return_value)
        children.extend(a_mock._mock_children.values())

        # Enqueue children.
        for child in children:
            # Create a new parent name for the child.
            if child._mock_new_name == '()':
                dot = ''
            else:
                dot = '.'
            if parent_name:
                new_name = parent_name + dot + child._mock_new_name
            else:
                new_name = child._mock_new_name

            queue.append((new_name, child))

        # Collect calls for the current mock.
        for call in a_mock._mock_await_args_list:
            new_name = parent_name + (call._mock_name or '')
            new_call = _Call((new_name, ) + call._get_call_arguments())
            awaits.append(new_call)

    result = []
    for call in mock.mock_calls:
        if call in awaits:
            awaits.remove(call)
            result.append(call)
    return result