async def _subscribe_channels(self, websocket_assistant: WSAssistant):
"""
Subscribes to order events and balance events.
:param ws: the websocket assistant used to connect to the exchange
"""
path_params = {'user': self._current_listen_key}
msg_subscribe_orders = stomper.subscribe(
CONSTANTS.ORDERS_STREAM.format(**path_params),
CONSTANTS.SUBSCRIPTION_ID_ORDERS,
ack="auto")
msg_subscribe_trades = stomper.subscribe(
CONSTANTS.TRADE_UPDATE_STREAM.format(**path_params),
CONSTANTS.SUBSCRIPTION_ID_TRADE_UPDATE,
ack="auto")
msg_subscribe_account = stomper.subscribe(
CONSTANTS.ACCOUNT_STREAM.format(**path_params),
CONSTANTS.SUBSCRIPTION_ID_ACCOUNT,
ack="auto")
_ = await safe_gather(
websocket_assistant.subscribe(request=WSPlainTextRequest(
payload=msg_subscribe_trades)),
websocket_assistant.subscribe(request=WSPlainTextRequest(
payload=msg_subscribe_orders)),
websocket_assistant.subscribe(request=WSPlainTextRequest(
payload=msg_subscribe_account)),
return_exceptions=True)
def test_ws_assistant_authenticates(self, send_mock):
class Auth(AuthBase):
async def rest_authenticate(self, request: RESTRequest) -> RESTRequest:
pass
async def ws_authenticate(self, request: WSRequest) -> WSRequest:
request.payload["authenticated"] = True
return request
ws_assistant = WSAssistant(connection=self.ws_connection, auth=Auth())
sent_requests = []
send_mock.side_effect = lambda r: sent_requests.append(r)
payload = {"one": 1}
req = WSRequest(payload)
auth_req = WSRequest(payload, is_auth_required=True)
self.async_run_with_timeout(ws_assistant.send(req))
self.async_run_with_timeout(ws_assistant.send(auth_req))
sent_request = sent_requests[0]
auth_sent_request = sent_requests[1]
expected = {"one": 1}
auth_expected = {"one": 1, "authenticated": True}
self.assertEqual(expected, sent_request.payload)
self.assertEqual(auth_expected, auth_sent_request.payload)
async def _process_ws_messages(self, websocket_assistant: WSAssistant, output: asyncio.Queue):
async for ws_response in websocket_assistant.iter_messages():
data = ws_response.data
if (data.get("type") == "message"
and data.get("subject") in [CONSTANTS.ORDER_CHANGE_EVENT_TYPE,
CONSTANTS.BALANCE_EVENT_TYPE]):
output.put_nowait(data)
async def _process_websocket_messages(self,
websocket_assistant: WSAssistant,
queue: asyncio.Queue):
async for ws_response in websocket_assistant.iter_messages():
data: Dict[str, Any] = ws_response.data
decompressed_data = utils.decompress_ws_message(data)
try:
if type(decompressed_data) == str:
json_data = json.loads(decompressed_data)
else:
json_data = decompressed_data
except asyncio.CancelledError:
raise
except Exception:
self.logger().warning(
f"Invalid event message received through the order book data source "
f"connection ({decompressed_data})")
continue
if "errorCode" in json_data or "errorMessage" in json_data:
raise ValueError(
f"Error message received in the order book data source: {json_data}"
)
await self._process_event_message(event_message=json_data,
queue=queue)
def test_receive_post_processes(self, receive_mock):
class SomePostProcessor(WSPostProcessorBase):
async def post_process(self, response_: WSResponse) -> WSResponse:
response_.data["two"] = 2
return response_
ws_assistant = WSAssistant(connection=self.ws_connection,
ws_post_processors=[SomePostProcessor()])
data = {"one": 1}
response_mock = WSResponse(data)
receive_mock.return_value = response_mock
response = self.async_run_with_timeout(ws_assistant.receive())
expected = {"one": 1, "two": 2}
self.assertEqual(expected, response.data)
async def _process_ws_messages(self, websocket_assistant: WSAssistant):
async for ws_response in websocket_assistant.iter_messages():
data = ws_response.data
if data.get("type") == "message":
event_type = data.get("subject")
if event_type in [
CONSTANTS.DIFF_EVENT_TYPE, CONSTANTS.TRADE_EVENT_TYPE
]:
self._message_queue[event_type].put_nowait(data)
async def _process_ws_messages(self, ws: WSAssistant, output: asyncio.Queue):
async for ws_response in ws.iter_messages():
data = ws_response.data
if isinstance(data, list):
for message in data:
if message["e"] in ["executionReport", "outboundAccountInfo"]:
output.put_nowait(message)
elif data.get("auth") == "fail":
raise IOError("Private channel authentication failed.")
def test_send_pre_processes(self, send_mock):
class SomePreProcessor(WSPreProcessorBase):
async def pre_process(self, request_: WSRequest) -> WSRequest:
request_.payload["two"] = 2
return request_
ws_assistant = WSAssistant(connection=self.ws_connection,
ws_pre_processors=[SomePreProcessor()])
sent_requests = []
send_mock.side_effect = lambda r: sent_requests.append(r)
payload = {"one": 1}
request = WSRequest(payload)
self.async_run_with_timeout(ws_assistant.send(request))
sent_request = sent_requests[0]
expected = {"one": 1, "two": 2}
self.assertEqual(expected, sent_request.payload)
async def _subscribe_channels(self, client: WSAssistant):
"""
Subscribes to the trade events and diff orders events through the provided websocket connection.
:param client: the websocket assistant used to connect to the exchange
"""
try:
subscriptions = []
for trading_pair in self._trading_pairs:
symbol = await self._connector.exchange_symbol_associated_to_pair(
trading_pair=trading_pair)
path_params = {'symbol': symbol}
msg_subscribe_books = stomper.subscribe(
CONSTANTS.BOOK_STREAM.format(**path_params),
f"{CONSTANTS.SUBSCRIPTION_ID_BOOKS}_{trading_pair}",
ack="auto")
msg_subscribe_trades = stomper.subscribe(
CONSTANTS.TRADES_STREAM.format(**path_params),
f"{CONSTANTS.SUBSCRIPTION_ID_TRADES}_{trading_pair}",
ack="auto")
subscriptions.append(
client.subscribe(
WSPlainTextRequest(payload=msg_subscribe_books)))
subscriptions.append(
client.subscribe(
WSPlainTextRequest(payload=msg_subscribe_trades)))
_ = await safe_gather(*subscriptions)
self.logger().info(
"Subscribed to public order book and trade channels...")
except asyncio.CancelledError:
raise
except Exception:
self.logger().error(
"Unexpected error occurred subscribing to order book trading and delta streams...",
exc_info=True)
raise
async def _iter_messages(self, ws: WSAssistant) -> AsyncIterable[Dict]:
"""
Generator function that returns messages from the web socket stream
:param ws: current web socket connection
:returns: message in AsyncIterable format
"""
# Terminate the recv() loop as soon as the next message timed out, so the outer loop can reconnect.
try:
async for response in ws.iter_messages():
msg = response.data
yield msg
except asyncio.TimeoutError:
self.logger().warning("WebSocket ping timed out. Going to reconnect...")
finally:
await ws.disconnect()
async def _process_ws_messages(self, ws: WSAssistant):
async for ws_response in ws.iter_messages():
data = ws_response.data
if data.get("msg") == "Success":
continue
event_type = data.get("topic")
if event_type == CONSTANTS.DIFF_EVENT_TYPE:
if data.get("f"):
self._message_queue[
CONSTANTS.SNAPSHOT_EVENT_TYPE].put_nowait(data)
else:
self._message_queue[CONSTANTS.DIFF_EVENT_TYPE].put_nowait(
data)
elif event_type == CONSTANTS.TRADE_EVENT_TYPE:
self._message_queue[CONSTANTS.TRADE_EVENT_TYPE].put_nowait(
data)
def setUp(self) -> None:
super().setUp()
aiohttp_client_session = aiohttp.ClientSession()
self.ws_connection = WSConnection(aiohttp_client_session)
self.ws_assistant = WSAssistant(self.ws_connection)
class WSAssistantTest(unittest.TestCase):
@classmethod
def setUpClass(cls) -> None:
super().setUpClass()
cls.ev_loop = asyncio.get_event_loop()
def setUp(self) -> None:
super().setUp()
aiohttp_client_session = aiohttp.ClientSession()
self.ws_connection = WSConnection(aiohttp_client_session)
self.ws_assistant = WSAssistant(self.ws_connection)
def async_run_with_timeout(self, coroutine: Awaitable, timeout: int = 1):
ret = self.ev_loop.run_until_complete(asyncio.wait_for(coroutine, timeout))
return ret
@patch("hummingbot.core.web_assistant.connections.ws_connection.WSConnection.connect")
def test_connect(self, connect_mock):
ws_url = "ws://some.url"
ping_timeout = 10
message_timeout = 20
self.async_run_with_timeout(
self.ws_assistant.connect(ws_url, ping_timeout=ping_timeout, message_timeout=message_timeout)
)
connect_mock.assert_called_with(ws_url, ping_timeout, message_timeout)
@patch("hummingbot.core.web_assistant.connections.ws_connection.WSConnection.disconnect")
def test_disconnect(self, disconnect_mock):
self.async_run_with_timeout(self.ws_assistant.disconnect())
disconnect_mock.assert_called()
@patch("hummingbot.core.web_assistant.connections.ws_connection.WSConnection.send")
def test_send(self, send_mock):
sent_requests = []
send_mock.side_effect = lambda r: sent_requests.append(r)
payload = {"one": 1}
request = WSRequest(payload)
self.async_run_with_timeout(self.ws_assistant.send(request))
self.assertEqual(1, len(sent_requests))
sent_request = sent_requests[0]
self.assertNotEqual(id(request), id(sent_request)) # has been cloned
self.assertEqual(request, sent_request)
@patch("hummingbot.core.web_assistant.connections.ws_connection.WSConnection.send")
def test_send_pre_processes(self, send_mock):
class SomePreProcessor(WSPreProcessorBase):
async def pre_process(self, request_: WSRequest) -> WSRequest:
request_.payload["two"] = 2
return request_
ws_assistant = WSAssistant(
connection=self.ws_connection, ws_pre_processors=[SomePreProcessor()]
)
sent_requests = []
send_mock.side_effect = lambda r: sent_requests.append(r)
payload = {"one": 1}
request = WSRequest(payload)
self.async_run_with_timeout(ws_assistant.send(request))
sent_request = sent_requests[0]
expected = {"one": 1, "two": 2}
self.assertEqual(expected, sent_request.payload)
@patch("hummingbot.core.web_assistant.connections.ws_connection.WSConnection.send")
def test_subscribe(self, send_mock):
sent_requests = []
send_mock.side_effect = lambda r: sent_requests.append(r)
payload = {"one": 1}
request = WSRequest(payload)
self.async_run_with_timeout(self.ws_assistant.subscribe(request))
self.assertEqual(1, len(sent_requests))
sent_request = sent_requests[0]
self.assertNotEqual(id(request), id(sent_request)) # has been cloned
self.assertEqual(request, sent_request)
@patch("hummingbot.core.web_assistant.connections.ws_connection.WSConnection.send")
def test_ws_assistant_authenticates(self, send_mock):
class Auth(AuthBase):
async def rest_authenticate(self, request: RESTRequest) -> RESTRequest:
pass
async def ws_authenticate(self, request: WSRequest) -> WSRequest:
request.payload["authenticated"] = True
return request
ws_assistant = WSAssistant(connection=self.ws_connection, auth=Auth())
sent_requests = []
send_mock.side_effect = lambda r: sent_requests.append(r)
payload = {"one": 1}
req = WSRequest(payload)
auth_req = WSRequest(payload, is_auth_required=True)
self.async_run_with_timeout(ws_assistant.send(req))
self.async_run_with_timeout(ws_assistant.send(auth_req))
sent_request = sent_requests[0]
auth_sent_request = sent_requests[1]
expected = {"one": 1}
auth_expected = {"one": 1, "authenticated": True}
self.assertEqual(expected, sent_request.payload)
self.assertEqual(auth_expected, auth_sent_request.payload)
@patch("hummingbot.core.web_assistant.connections.ws_connection.WSConnection.receive")
def test_receive(self, receive_mock):
data = {"one": 1}
response_mock = WSResponse(data)
receive_mock.return_value = response_mock
response = self.async_run_with_timeout(self.ws_assistant.receive())
self.assertEqual(data, response.data)
@patch("hummingbot.core.web_assistant.connections.ws_connection.WSConnection.receive")
def test_receive_post_processes(self, receive_mock):
class SomePostProcessor(WSPostProcessorBase):
async def post_process(self, response_: WSResponse) -> WSResponse:
response_.data["two"] = 2
return response_
ws_assistant = WSAssistant(
connection=self.ws_connection, ws_post_processors=[SomePostProcessor()]
)
data = {"one": 1}
response_mock = WSResponse(data)
receive_mock.return_value = response_mock
response = self.async_run_with_timeout(ws_assistant.receive())
expected = {"one": 1, "two": 2}
self.assertEqual(expected, response.data)
@patch(
"hummingbot.core.web_assistant.connections.ws_connection.WSConnection.connected",
new_callable=PropertyMock,
)
@patch("hummingbot.core.web_assistant.connections.ws_connection.WSConnection.receive")
def test_iter_messages(self, receive_mock, connected_mock):
connected_mock.return_value = True
data = {"one": 1}
response_mock = WSResponse(data)
receive_mock.return_value = response_mock
iter_messages_iterator = self.ws_assistant.iter_messages()
response = self.async_run_with_timeout(iter_messages_iterator.__anext__())
self.assertEqual(data, response.data)
connected_mock.return_value = False
with self.assertRaises(StopAsyncIteration):
self.async_run_with_timeout(iter_messages_iterator.__anext__())
async def get_ws_assistant(self) -> WSAssistant:
connection = await self._connections_factory.get_ws_connection()
assistant = WSAssistant(connection, self._ws_pre_processors,
self._ws_post_processors, self._auth)
return assistant
async def _process_websocket_messages(self, websocket_assistant: WSAssistant, queue: asyncio.Queue):
async for ws_response in websocket_assistant.iter_messages():
data = ws_response.data
await self._process_event_message(event_message=data, queue=queue)