async def test_check_connection(ds: IdexDatastream):
ds.init = CoroutineMock()
ds._ws = object
await ds._check_connection()
ds.init.assert_not_awaited()
ds._ws = None
await ds._check_connection()
ds.init.assert_awaited_once()
async def test_listen_reconnect(ds: IdexDatastream):
class BreakExc(Exception):
pass
exc = websockets.ConnectionClosed(code=999, reason='some reason')
exit_exc = BreakExc('to break the infinite loop')
ds._check_connection = CoroutineMock()
ds._ws = MagicMock()
ds._ws.closed = False
ds._ws.__aiter__.side_effect = exc
ds.init = CoroutineMock()
ds.sub_manager.resubscribe = CoroutineMock()
ds.sub_manager.resubscribe.side_effect = exit_exc
ds._logger.error = Mock()
with pytest.raises(BreakExc):
async for m in ds.listen():
break
ds._check_connection.assert_awaited_once()
ds._logger.error.assert_called_once_with(exc)
ds.init.assert_awaited_once()
ds.sub_manager.resubscribe.assert_awaited_once()
async def test_listen_raise(ds: IdexDatastream):
class UnhandledExc(Exception):
pass
exc = UnhandledExc('some unknown exception')
ds._check_connection = CoroutineMock()
ds._ws = MagicMock()
ds._ws.closed = False
ds._ws.__aiter__.side_effect = exc
with pytest.raises(UnhandledExc):
async for m in ds.listen():
break
ds._check_connection.assert_awaited_once()
async def test_create_connection(ds: IdexDatastream):
ds._check_connection = CoroutineMock()
ds._get_rid = Mock(return_value='rid:smth')
ds._compose_message = Mock(return_value={'some': 'data'})
ds._ws = Mock()
ds._ws.send = CoroutineMock()
ds._encode = Mock()
result = await ds.send_message('some_request', {'some': 'payload'})
ds._check_connection.assert_awaited_once()
ds._get_rid.assert_called_once()
ds._compose_message.assert_called_once_with('rid:smth', 'some_request',
{'some': 'payload'})
ds._encode.assert_called_once_with({'some': 'data'})
ds._ws.send.assert_awaited_once()
assert result == 'rid:smth'
async def test_wait_for_handshake_response(ds: IdexDatastream):
ds._ws = CoroutineMock()
loop_mock = CoroutineMock()
ds._loop = loop_mock
recv_mock = Mock()
recv_mock.return_value = None
ds._ws.recv = recv_mock
with patch('asyncio.wait_for', new=CoroutineMock()) as mock:
mock.return_value = '{"some":"data"}'
await ds._wait_for_handshake_response()
mock.assert_awaited_once_with(None,
ds._HANDSHAKE_TIMEOUT,
loop=loop_mock)
async def test_listen(ds: IdexDatastream):
msg_data = '{"payload":"{"some": "data"}"}'
ds._check_connection = CoroutineMock()
ds._ws = MagicMock()
ds._ws.closed = False
ds._ws.__aiter__.return_value = (msg_data, )
processed_message = {'a': 'b'}
ds._process_message: Mock = Mock()
ds._process_message.return_value = processed_message
msg = None
async for m in ds.listen():
msg = m
break
ds._process_message.assert_called_once_with(msg_data)
assert msg == processed_message