Python asyncio: how to mock __aiter__() method?

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)

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.

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]

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