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.
With Mock you can mock magic methods but you have to define them. MagicMock has "default implementations of most of the magic methods.". If you don't need to test any magic methods, Mock is adequate and doesn't bring a lot of extraneous things into your tests.
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.
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
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)
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With