Python asyncio: how to mock __aiter__() method?

Question

I have a code which is listening to messages on WebSocket using aiohttp.

It looks like:

async for msg in ws:
    await self._ws_msg_handler.handle_message(ws, msg, _services)

Where ws is an instance of aiohttp.web.WebSocketResponse() (original code)

In my test I mock WebSocketResponse() and its __aiter__ method:

def coro_mock(**kwargs):
    return asyncio.coroutine(mock.Mock(**kwargs))


@pytest.mark.asyncio
@mock.patch('aiojsonrpc.request_handler.WebSocketMessageHandler')
async def test_rpc_websocket_handler(
    MockWebSocketMessageHandler,
    rpc_websocket_handler
):

    ws_response = 'aiojsonrpc.request_handler.WebSocketResponse'
    with mock.patch(ws_response) as MockWebSocketResponse:
        MockRequest = mock.MagicMock()
        req = MockRequest()

        ws_instance = MockWebSocketResponse.return_value
        ws_instance.prepare = coro_mock()
        ws_instance.__aiter__ = coro_mock(return_value=iter(range(5)))
        ws_instance.__anext__ = coro_mock()

        handle_msg_result = 'Message processed'
        MockWebSocketMessageHandler.handle_message.side_effect = Exception(
            handle_msg_result)
        msg_handler = MockWebSocketMessageHandler()

        with pytest.raises(Exception) as e:
            await request_handler.RpcWebsocketHandler(msg_handler)(req)
        assert str(e.value) == handle_msg_result

Though when I run the test it fails with the error message saying:

'async for' requires an object with __aiter__ method, got MagicMock

=================================================================================== FAILURES ===================================================================================
__________________________________________________________________________ test_rpc_websocket_handler __________________________________________________________________________

MockWebSocketMessageHandler = <MagicMock name='WebSocketMessageHandler' id='140687969989632'>
rpc_websocket_handler = <aiojsonrpc.request_handler.RpcWebsocketHandler object at 0x7ff47879b0f0>

    @pytest.mark.asyncio
    @mock.patch('aiojsonrpc.request_handler.WebSocketMessageHandler')
    async def test_rpc_websocket_handler(
        MockWebSocketMessageHandler,
        rpc_websocket_handler
    ):

        ws_response = 'aiojsonrpc.request_handler.WebSocketResponse'
        with mock.patch(ws_response) as MockWebSocketResponse:
            # MockRequest = mock.create_autospec(aiohttp.web_reqrep.Request)
            # req = MockRequest(*[None] * 6)
            MockRequest = mock.MagicMock()
            req = MockRequest()

            ws_instance = MockWebSocketResponse.return_value
            ret = mock.Mock()
            ws_instance.prepare = coro_mock()
            ws_instance.__aiter__ = coro_mock(return_value=iter(range(5)))
            ws_instance.__anext__ = coro_mock()

            handle_msg_result = 'Message processed'
            MockWebSocketMessageHandler.handle_message.side_effect = Exception(
                handle_msg_result)
            msg_handler = MockWebSocketMessageHandler()

            with pytest.raises(Exception) as e:
                await request_handler.RpcWebsocketHandler(msg_handler)(req)
>           assert str(e.value) == handle_msg_result
E           assert "'async for' ...got MagicMock" == 'Message processed'
E             - 'async for' requires an object with __aiter__ method, got MagicMock
E             + Message processed

tests/test_request_handler.py:252: AssertionError

So it behaves like __aiter__() was never mocked. How I'm supposed to accomplish correct mocking in this case?

---

Update:

For now I've found a workaround to make the code testable though I would really appreciate if someone tell me how to deal with the issue described in the original question.

Answer 1

I have a python version that supports AsyncMock and I also leverage pytest_mock. I came up with this solution to this problem combining the use of AsyncMock side_effect:

from typing import List
import pytest
import asyncio

from pytest_mock.plugin import MockerFixture


pytestmark = pytest.mark.asyncio


async def async_generator(numbers: List[int]):
    for number in numbers:
        yield number
        await asyncio.sleep(0.1)


async def function_to_test(numbers: List[int]):
    async for thing in async_generator(numbers):
        yield thing * 3
        await asyncio.sleep(0.1)


async def test_async_generator(mocker: MockerFixture):
    mock_numbers = [1, 2, 3, 4, 5]

    async def async_generator_side_effect(numbers: List[int]):
        for number in numbers:
            yield number

    mock_async_generator = mocker.patch("tests.test_async_generator.async_generator")
    mock_async_generator.side_effect = async_generator_side_effect

    actual = []
    async for result in function_to_test(mock_numbers):
        actual.append(result)

    assert actual == [3, 6, 9, 12, 15]

Answer 2

Works for py38

from unittest.mock import MagicMock

async def test_iterable(self):
    loop_iterations = 0
    mock = MagicMock()
    mock.__aiter__.return_value = range(5)
    async for _ in mock:
        loop_iterations += 1

    self.assertEqual(5, loop_iterations)

Answer 3

You can make the mocked class return an object implementing the expected interface:

class AsyncIterator:
    def __init__(self, seq):
        self.iter = iter(seq)

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            return next(self.iter)
        except StopIteration:
            raise StopAsyncIteration

MockWebSocketResponse.return_value = AsyncIterator(range(5))

I don't think there is a way (yet) to correctly mock an object implementing __aiter__, it may be a python bug, as async for rejects a MagicMock, even if hasattr(the_magic_mock, '__aiter__') is True.

EDIT (13/12/2017): the library asynctest supports asynchronous iterators and context managers since 0.11, asynctest.MagicMock provides this feature for free.

Answer 4

For posterity, I had the same problem of needing to test an async for loop, but the accepted solution doesn't seem to work for Python 3.7. The example below works for 3.6.x and 3.7.0, but not for 3.5.x:

import asyncio


class AsyncIter:    
    def __init__(self, items):    
        self.items = items    

    async def __aiter__(self):    
        for item in self.items:    
            yield item    


async def print_iter(items):
    async for item in items:
        print(item)


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    things = AsyncIter([1, 2, 3])
    loop.run_until_complete(print_iter(things))
    loop.close()

With the above, mocking it looks something like:

with mock.patch('some.async.iter', return_value=AsyncIter([1, 2, 3])):
  # do test requiring mocked iter