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 _subscribe_channels(self, ws: WSAssistant):
try:
for pair in self._trading_pairs:
subscribe_orderbook_request: WSRequest = WSRequest({
"type":
"subscribe",
"channel":
self.ORDERBOOK_CHANNEL,
"id":
pair,
})
subscribe_trade_request: WSRequest = WSRequest({
"type":
"subscribe",
"channel":
self.TRADE_CHANNEL,
"id":
pair,
})
await ws.send(subscribe_orderbook_request)
await ws.send(subscribe_trade_request)
self.logger().info(
"Subscribed to public orderbook 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 _subscribe_channels(self, ws: WSAssistant):
"""
Subscribes to the trade events and diff orders events through the provided websocket connection.
:param ws: the websocket assistant used to connect to the exchange
"""
try:
trade_params = []
depth_params = []
for trading_pair in self._trading_pairs:
symbol = await self.exchange_symbol_associated_to_pair(
trading_pair=trading_pair,
domain=self._domain,
api_factory=self._api_factory,
throttler=self._throttler,
time_synchronizer=self._time_synchronizer)
trade_params.append(f"{symbol.lower()}@trade")
depth_params.append(f"{symbol.lower()}@depth@100ms")
payload = {"method": "SUBSCRIBE", "params": trade_params, "id": 1}
subscribe_trade_request: WSRequest = WSRequest(payload=payload)
payload = {"method": "SUBSCRIBE", "params": depth_params, "id": 2}
subscribe_orderbook_request: WSRequest = WSRequest(payload=payload)
await ws.send(subscribe_trade_request)
await ws.send(subscribe_orderbook_request)
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 ws_auth(self) -> Dict[Any, Any]:
"""ws private (user account balance request)"""
listen_key = (
await self.rest_auth(CONSTANTS.USER_ID_PATH_URL)
)["id"] # this is the getUserId from the github latoken api client library
client: WSAssistant = await self.api_factory.get_ws_assistant()
await client.connect(ws_url=web_utils.ws_url(self.domain),
ping_timeout=CONSTANTS.WS_HEARTBEAT_TIME_INTERVAL)
msg_out = stomper.Frame()
msg_out.cmd = "CONNECT"
msg_out.headers.update({"accept-version": "1.1", "heart-beat": "0,0"})
connect_request: WSRequest = WSRequest(payload=msg_out.pack(),
is_auth_required=True)
await client.send(connect_request)
await client.receive()
path_params = {'user': listen_key}
msg_subscribe_account = stomper.subscribe(
CONSTANTS.ACCOUNT_STREAM.format(**path_params),
CONSTANTS.SUBSCRIPTION_ID_ACCOUNT,
ack="auto")
_ = await client.subscribe(request=WSRequest(
payload=msg_subscribe_account))
response = []
async for ws_response in client.iter_messages():
msg_in = stomper.Frame()
data = msg_in.unpack(ws_response.data.decode())
event_type = int(data['headers']['subscription'].split('_')[0])
if event_type == CONSTANTS.SUBSCRIPTION_ID_ACCOUNT:
response.append(ujson.loads(data["body"])["payload"])
break
await client.disconnect()
return response[0]
async def _subscribe_channels(self, ws: WSAssistant):
"""
Subscribes to order events and balance events.
:param ws: the websocket assistant used to connect to the exchange
"""
try:
orders_change_payload = {
"id": web_utils.next_message_id(),
"type": "subscribe",
"topic": "/spotMarket/tradeOrders",
"privateChannel": True,
"response": False,
}
subscribe_order_change_request: WSRequest = WSRequest(payload=orders_change_payload)
balance_payload = {
"id": web_utils.next_message_id(),
"type": "subscribe",
"topic": "/account/balance",
"privateChannel": True,
"response": False,
}
subscribe_balance_request: WSRequest = WSRequest(payload=balance_payload)
await ws.send(subscribe_order_change_request)
await ws.send(subscribe_balance_request)
self.logger().info("Subscribed to private order changes and balance updates channels...")
except asyncio.CancelledError:
raise
except Exception:
self.logger().exception("Unexpected error occurred subscribing to user streams...")
raise
async def _subscribe_channels(self, ws: WSAssistant):
"""
Subscribes to the trade events and diff orders events through the provided websocket connection.
:param ws: the websocket assistant used to connect to the exchange
"""
try:
symbols = ",".join([
await self.exchange_symbol_associated_to_pair(
trading_pair=pair,
domain=self._domain,
api_factory=self._api_factory,
throttler=self._throttler,
time_synchronizer=self._time_synchronizer)
for pair in self._trading_pairs
])
trades_payload = {
"id": web_utils.next_message_id(),
"type": "subscribe",
"topic": f"/market/match:{symbols}",
"privateChannel": False,
"response": False,
}
subscribe_trade_request: WSRequest = WSRequest(
payload=trades_payload)
order_book_payload = {
"id": web_utils.next_message_id(),
"type": "subscribe",
"topic": f"/market/level2:{symbols}",
"privateChannel": False,
"response": False,
}
subscribe_orderbook_request: WSRequest = WSRequest(
payload=order_book_payload)
await ws.send(subscribe_trade_request)
await ws.send(subscribe_orderbook_request)
self.logger().info(
"Subscribed to public order book and trade channels...")
except asyncio.CancelledError:
raise
except Exception:
self.logger().exception(
"Unexpected error occurred subscribing to order book trading and delta streams..."
)
raise
async def _subscribe_channels(self, ws: WSAssistant):
"""
Subscribes to the trade events and diff orders events through the provided websocket connection.
:param ws: the websocket assistant used to connect to the exchange
"""
try:
trade_params = []
depth_params = []
for trading_pair in self._trading_pairs:
symbol = await self.exchange_symbol_associated_to_pair(
trading_pair=trading_pair,
domain=self._domain,
api_factory=self._api_factory,
throttler=self._throttler)
trade_params.append(f"trade:{symbol}")
depth_params.append(f"depth:{symbol}")
payload: Dict[str, str] = {
"op": "subscribe",
"args": trade_params + depth_params,
}
subscribe_request: WSRequest = WSRequest(payload=payload)
await ws.send(subscribe_request)
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 listen_for_user_stream(self, ev_loop: asyncio.BaseEventLoop,
output: asyncio.Queue):
ws = None
while True:
try:
self.logger().info(f"Connecting to {CONSTANTS.DYDX_WS_URL}")
ws: WSAssistant = await self._get_ws_assistant()
await ws.connect(ws_url=CONSTANTS.DYDX_WS_URL,
ping_timeout=self.HEARTBEAT_INTERVAL)
auth_params = self._dydx_auth.get_ws_auth_params()
auth_request: WSRequest = WSRequest(auth_params)
await ws.send(auth_request)
self.logger().info("Authenticated user stream...")
async for ws_response in ws.iter_messages():
data = ws_response.data
if data.get("type", "") in ["subscribed", "channel_data"]:
output.put_nowait(data)
except asyncio.CancelledError:
raise
except Exception:
self.logger().error(
"Unexpected error with dydx WebSocket connection. "
"Retrying after 30 seconds...",
exc_info=True)
finally:
# Make sure no background tasks is leaked
ws and await ws.disconnect()
await self._sleep(30.0)
async def _authenticate(self, ws: WSAssistant):
"""
Authenticates user to websocket
"""
try:
auth_payload: Dict[str,
Any] = self._bitmart_auth.get_ws_auth_payload(
bitmart_utils.get_ms_timestamp())
ws_message: WSRequest = WSRequest(auth_payload)
await ws.send(ws_message)
ws_response = await ws.receive()
auth_resp: Dict[str, Any] = ws_response.data
if "errorCode" in auth_resp.keys():
self.logger().error(
f"WebSocket login errored with message: {auth_resp['errorMessage']}",
exc_info=True)
raise ConnectionError
except asyncio.CancelledError:
raise
except Exception:
self.logger().error(
"Error occurred when authenticating to user stream.",
exc_info=True)
raise
async def listen_for_user_stream(self, output: asyncio.Queue):
while True:
try:
# Initialize Websocket Connection
self.logger().info(f"Connecting to {CONSTANTS.WS_PRIVATE_URL}")
await self._get_ws_assistant()
await self._ws_assistant.connect(
ws_url=CONSTANTS.WS_PRIVATE_URL,
ping_timeout=self.HEARTBEAT_INTERVAL)
await self._authenticate_client()
await self._subscribe_channels()
async for ws_response in self._ws_assistant.iter_messages():
data = ws_response.data
if data["action"] == "ping":
pong_request = WSRequest(payload={
"action": "pong",
"data": data["data"]
})
await self._ws_assistant.send(request=pong_request)
continue
output.put_nowait(data)
except asyncio.CancelledError:
raise
except Exception:
self.logger().error(
"Unexpected error with Huobi WebSocket connection. "
"Retrying after 30 seconds...",
exc_info=True)
finally:
self._ws_assistant and await self._ws_assistant.disconnect()
await self._sleep(30.0)
def test_no_auth_added_to_wsrequest(self):
payload = {"param1": "value_param_1"}
request = WSRequest(payload=payload, is_auth_required=True)
self.async_run_with_timeout(self.auth.ws_authenticate(request))
self.assertEqual(payload, request.payload)
def test_send_when_disconnected_raises(self):
request = WSRequest(payload={"one": 1})
with self.assertRaises(RuntimeError) as e:
self.async_run_with_timeout(self.ws_connection.send(request))
self.assertEqual("WS is not connected.", str(e.exception))
async def _subscribe_to_order_book_streams(self) -> WSAssistant:
url = f"{utils.wss_url(CONSTANTS.PUBLIC_WS_ENDPOINT, self._domain)}"
ws: WSAssistant = await self._get_ws_assistant()
await ws.connect(ws_url=url, ping_timeout=CONSTANTS.HEARTBEAT_TIME_INTERVAL)
stream_id_channel_pairs = [
(CONSTANTS.DIFF_STREAM_ID, "@depth"),
(CONSTANTS.TRADE_STREAM_ID, "@aggTrade"),
(CONSTANTS.FUNDING_INFO_STREAM_ID, "@markPrice"),
]
for stream_id, channel in stream_id_channel_pairs:
params = []
for trading_pair in self._trading_pairs:
symbol = await self.convert_to_exchange_trading_pair(
hb_trading_pair=trading_pair,
domain=self._domain,
throttler=self._throttler)
params.append(f"{symbol.lower()}{channel}")
payload = {
"method": "SUBSCRIBE",
"params": params,
"id": stream_id,
}
subscribe_request: WSRequest = WSRequest(payload)
await ws.send(subscribe_request)
return ws
async def listen_for_user_stream(self, output: asyncio.Queue):
"""
*required
Subscribe to user stream via web socket, and keep the connection open for incoming messages
:param output: an async queue where the incoming messages are stored
"""
ws = None
while True:
try:
rest_assistant = await self._api_factory.get_rest_assistant()
url = f"{CONSTANTS.REST_URL}/{CONSTANTS.INFO_PATH_URL}"
request = RESTRequest(method=RESTMethod.GET,
url=url,
endpoint_url=CONSTANTS.INFO_PATH_URL,
is_auth_required=True)
async with self._throttler.execute_task(
CONSTANTS.INFO_PATH_URL):
response: RESTResponse = await rest_assistant.call(
request=request)
info = await response.json()
accountGroup = info.get("data").get("accountGroup")
headers = self._ascend_ex_auth.get_auth_headers(
CONSTANTS.STREAM_PATH_URL)
payload = {
"op": CONSTANTS.SUB_ENDPOINT_NAME,
"ch": "order:cash"
}
ws: WSAssistant = await self._get_ws_assistant()
url = f"{get_ws_url_private(accountGroup)}/{CONSTANTS.STREAM_PATH_URL}"
await ws.connect(ws_url=url,
ws_headers=headers,
ping_timeout=self.HEARTBEAT_PING_INTERVAL)
subscribe_request: WSRequest = WSRequest(payload)
async with self._throttler.execute_task(
CONSTANTS.SUB_ENDPOINT_NAME):
await ws.send(subscribe_request)
async for raw_msg in ws.iter_messages():
msg = raw_msg.data
if msg is None:
continue
self._last_recv_time = time.time()
output.put_nowait(msg)
except asyncio.CancelledError:
raise
except Exception:
self.logger().error(
"Unexpected error with AscendEx WebSocket connection. "
"Retrying after 30 seconds...",
exc_info=True)
await self._sleep(30.0)
finally:
ws and await ws.disconnect()
async def _handle_ping_message(self, ws: aiohttp.ClientWebSocketResponse):
"""
Responds with pong to a ping message send by a server to keep a websocket connection alive
"""
async with self._throttler.execute_task(CONSTANTS.PONG_ENDPOINT_NAME):
payload = {"op": "pong"}
pong_request: WSRequest = WSRequest(payload)
await ws.send(pong_request)
def test_ws_authenticate(self):
request: WSRequest = WSRequest(payload={"TEST": "SOME_TEST_PAYLOAD"},
throttler_limit_id="TEST_LIMIT_ID",
is_auth_required=True)
signed_request: WSRequest = self.async_run_with_timeout(
self.auth.ws_authenticate(request))
self.assertEqual(request, signed_request)
async def listen_for_subscriptions(self):
"""
Connects to the trade events and order diffs websocket endpoints and listens to the messages sent by the
exchange. Each message is stored in its own queue.
"""
ws = None
while True:
try:
connection_info = await web_utils.api_request(
path=CONSTANTS.PUBLIC_WS_DATA_PATH_URL,
api_factory=self._api_factory,
throttler=self._throttler,
domain=self._domain,
method=RESTMethod.POST,
)
ws_url = connection_info["data"]["instanceServers"][0][
"endpoint"]
ping_interval = int(
int(connection_info["data"]["instanceServers"][0]
["pingInterval"]) * 0.8 * 1e-3)
token = connection_info["data"]["token"]
ws: WSAssistant = await self._api_factory.get_ws_assistant()
await ws.connect(ws_url=f"{ws_url}?token={token}",
message_timeout=ping_interval)
await self._subscribe_channels(ws)
self._last_ws_message_sent_timestamp = self._time()
while True:
try:
seconds_until_next_ping = ping_interval - (
self._time() -
self._last_ws_message_sent_timestamp)
await asyncio.wait_for(
self._process_ws_messages(websocket_assistant=ws),
timeout=seconds_until_next_ping)
except asyncio.TimeoutError:
payload = {
"id": web_utils.next_message_id(),
"type": "ping",
}
ping_request = WSRequest(payload=payload)
self._last_ws_message_sent_timestamp = self._time()
await ws.send(request=ping_request)
except asyncio.CancelledError:
raise
except Exception:
self.logger().error(
"Unexpected error occurred when listening to order book streams. Retrying in 5 seconds...",
exc_info=True,
)
await self._sleep(5.0)
finally:
ws and await ws.disconnect()
async def ws_auth(self) -> Dict[Any, Any]:
ws = await self._api_factory.get_ws_assistant()
await ws.connect(
ws_url=coinflex_utils.websocket_url(domain=self._domain),
ping_timeout=CONSTANTS.WS_HEARTBEAT_TIME_INTERVAL)
async with timeout(30):
await ws.send(WSRequest({}, is_auth_required=True))
payload = {
"op": "subscribe",
"args": CONSTANTS.WS_CHANNELS["USER_STREAM"],
"tag": 101
}
await ws.send(WSRequest(payload=payload))
async for response in ws.iter_messages():
if all(x in response.data
for x in ["channel", "success", "event"]):
if response.data["channel"] == "balance:all":
return response.data
return None
async def pre_process(self, response: WSRequest) -> WSRequest:
if response.payload == CONSTANTS.WS_CONNECT_MSG:
msg_out = stomper.Frame()
msg_out.cmd = CONSTANTS.WS_CONNECT_MSG
msg_out.headers.update({
"accept-version": "1.1",
"heart-beat": "0,0"
})
response.payload = msg_out.pack()
return response
async def listen_for_user_stream(self, ev_loop: asyncio.AbstractEventLoop,
output: asyncio.Queue):
"""
*required
Subscribe to user stream via web socket, and keep the connection open for incoming messages
:param ev_loop: ev_loop to execute this function in
:param output: an async queue where the incoming messages are stored
"""
while True:
try:
self._ws_assistant = await self._web_assistants_factory.get_ws_assistant(
)
await self._ws_assistant.connect(
CONSTANTS.WS_URL,
message_timeout=CONSTANTS.WS_MESSAGE_TIMEOUT)
subscribe_payload: Dict[str, any] = {
"type": "subscribe",
"product_ids": self._trading_pairs,
"channels": [CONSTANTS.USER_CHANNEL_NAME]
}
subscribe_request = WSRequest(payload=subscribe_payload,
is_auth_required=True)
await self._ws_assistant.subscribe(subscribe_request)
async for msg in self._iter_messages(self._ws_assistant):
msg_type: str = msg.get("type", None)
if msg_type is None:
raise ValueError(
f"Coinbase Pro Websocket message does not contain a type - {msg}"
)
elif msg_type == "error":
raise ValueError(
f"Coinbase Pro Websocket received error message - {msg['message']}"
)
elif msg_type in ["open", "match", "change", "done"]:
output.put_nowait(msg)
elif msg_type in ["received", "activate", "subscriptions"]:
# these messages are not needed to track the order book
pass
else:
raise ValueError(
f"Unrecognized Coinbase Pro Websocket message received - {msg}"
)
except asyncio.CancelledError:
self._ws_assistant = None
raise
except Exception:
self._ws_assistant = None
self.logger().network(
"Unexpected error with WebSocket connection.",
exc_info=True,
app_warning_msg=
f"Unexpected error with WebSocket connection."
f" Retrying in {CONSTANTS.REST_API_LIMIT_COOLDOWN} seconds."
f" Check network connection.")
await self._sleep(CONSTANTS.REST_API_LIMIT_COOLDOWN)
async def listen_for_user_stream(self, output: asyncio.Queue):
"""
Connects to the user private channel in the exchange using a websocket connection. With the established
connection, listens to all balance events and order updates provided by the exchange and stores them in the
output queue
:param output: The queue where all received events should be stored
"""
ws: Optional[WSAssistant] = None
while True:
try:
connection_info = await web_utils.api_request(
path=CONSTANTS.PRIVATE_WS_DATA_PATH_URL,
api_factory=self._api_factory,
throttler=self._throttler,
domain=self._domain,
method=RESTMethod.POST,
is_auth_required=True,
)
ws_url = connection_info["data"]["instanceServers"][0]["endpoint"]
ping_interval = int(int(connection_info["data"]["instanceServers"][0]["pingInterval"]) * 0.8 * 1e-3)
token = connection_info["data"]["token"]
ws = await self._get_ws_assistant()
await ws.connect(ws_url=f"{ws_url}?token={token}", message_timeout=ping_interval)
await ws.ping() # to update last_recv_timestamp
await self._subscribe_channels(ws)
self._last_ws_message_sent_timestamp = self._time()
while True:
try:
seconds_until_next_ping = ping_interval - (self._time() - self._last_ws_message_sent_timestamp)
await asyncio.wait_for(self._process_ws_messages(websocket_assistant=ws, output=output),
timeout=seconds_until_next_ping)
except asyncio.TimeoutError:
payload = {
"id": web_utils.next_message_id(),
"type": "ping",
}
ping_request = WSRequest(payload=payload)
self._last_ws_message_sent_timestamp = self._time()
await ws.send(request=ping_request)
except asyncio.CancelledError:
raise
except Exception:
self.logger().exception(
"Unexpected error occurred when listening to user streams. Retrying in 5 seconds...")
await self._sleep(5.0)
finally:
ws and await ws.disconnect()
def test_send(self, ws_connect_mock):
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
self.async_run_with_timeout(self.ws_connection.connect(self.ws_url))
request = WSRequest(payload={"one": 1})
self.async_run_with_timeout(self.ws_connection.send(request))
json_msgs = self.mocking_assistant.json_messages_sent_through_websocket(
ws_connect_mock.return_value)
self.assertEqual(1, len(json_msgs))
self.assertEqual(request.payload, json_msgs[0])
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)
async def get_ws_subscription_message(self, subscription_type: str):
trading_pairs: List[str] = []
for tp in self._trading_pairs:
trading_pairs.append(convert_to_exchange_trading_pair(tp, '/'))
ws_message: WSRequest = WSRequest({
"event": "subscribe",
"pair": trading_pairs,
"subscription": {
"name": subscription_type,
"depth": 1000
}
})
return ws_message
async def listen_for_order_book_diffs(self, ev_loop: asyncio.AbstractEventLoop, output: asyncio.Queue):
"""
*required
Subscribe to diff channel via web socket, and keep the connection open for incoming messages
:param ev_loop: ev_loop to execute this function in
:param output: an async queue where the incoming messages are stored
"""
while True:
try:
trading_pairs: List[str] = self._trading_pairs
ws_assistant = await self._web_assistants_factory.get_ws_assistant()
await ws_assistant.connect(CONSTANTS.WS_URL, message_timeout=CONSTANTS.WS_MESSAGE_TIMEOUT)
subscribe_payload = {
"type": "subscribe",
"product_ids": trading_pairs,
"channels": [CONSTANTS.FULL_CHANNEL_NAME]
}
subscribe_request = WSRequest(payload=subscribe_payload)
await ws_assistant.subscribe(subscribe_request)
async for msg in self._iter_messages(ws_assistant):
msg_type: str = msg.get("type", None)
if msg_type is None:
raise ValueError(f"Coinbase Pro Websocket message does not contain a type - {msg}")
elif msg_type == "error":
raise ValueError(f"Coinbase Pro Websocket received error message - {msg['message']}")
elif msg_type in ["open", "match", "change", "done"]:
if msg_type == "done" and "price" not in msg:
# done messages with no price are completed market orders which can be ignored
continue
order_book_message: OrderBookMessage = CoinbaseProOrderBook.diff_message_from_exchange(msg)
output.put_nowait(order_book_message)
elif msg_type in ["received", "activate", "subscriptions"]:
# these messages are not needed to track the order book
continue
else:
raise ValueError(f"Unrecognized Coinbase Pro Websocket message received - {msg}")
except asyncio.CancelledError:
raise
except Exception:
self.logger().network(
"Unexpected error with WebSocket connection.",
exc_info=True,
app_warning_msg=f"Unexpected error with WebSocket connection."
f" Retrying in {CONSTANTS.REST_API_LIMIT_COOLDOWN} seconds."
f" Check network connection."
)
await self._sleep(CONSTANTS.WS_RECONNECT_COOLDOWN)
async def _subscribe_topic(self, topic: str):
try:
subscribe_request: WSRequest = WSRequest({
"action": "sub",
"ch": topic
})
await self._ws_assistant.send(subscribe_request)
resp: WSResponse = await self._ws_assistant.receive()
sub_response = resp.data
if sub_response.get("code", 0) != 200:
raise ValueError(f"Error subscribing to topic: {topic}")
self.logger().info(f"Successfully subscribed to {topic}")
except asyncio.CancelledError:
raise
except Exception:
self.logger().error(
f"Cannot subscribe to user stream topic: {topic}")
raise
async def listen_for_user_stream(self, output: asyncio.Queue):
ws = None
while True:
try:
async with self._throttler.execute_task(
CONSTANTS.WS_CONNECTION_LIMIT_ID):
ws: WSAssistant = await self._api_factory.get_ws_assistant(
)
await ws.connect(ws_url=CONSTANTS.WS_AUTH_URL,
ping_timeout=PING_TIMEOUT)
if self._current_auth_token is None:
self._current_auth_token = await self.get_auth_token()
for subscription_type in ["openOrders", "ownTrades"]:
subscribe_request: WSRequest = WSRequest({
"event": "subscribe",
"subscription": {
"name": subscription_type,
"token": self._current_auth_token
}
})
await ws.send(subscribe_request)
async for ws_response in ws.iter_messages():
msg = ws_response.data
if not (type(msg) is dict and "event" in msg.keys()
and msg["event"] in [
"heartbeat", "systemStatus",
"subscriptionStatus"
]):
output.put_nowait(msg)
except asyncio.CancelledError:
raise
except Exception:
self.logger().error(
"Unexpected error with Kraken WebSocket connection. "
"Retrying after 30 seconds...",
exc_info=True)
self._current_auth_token = None
await asyncio.sleep(30.0)
finally:
if ws is not None:
await ws.disconnect()
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 _emit(self,
channel: str,
data: Optional[Dict[str, Any]] = None) -> int:
data = data or {}
payload = {
"time": int(time.time()),
"channel": channel,
**data,
}
# if auth class was passed into websocket class
# we need to emit authenticated requests
if self._is_private:
payload["auth"] = self._auth.generate_auth_dict_ws(payload)
request = WSRequest(payload)
await self._ws_assistant.send(request)
return payload["time"]
async def ws_authenticate(self, request: WSRequest) -> WSRequest:
"""
This method is intended to configure a websocket request to be authenticated.
It should be used with empty requests to send an initial login payload.
:param request: the request to be configured for authenticated interaction
"""
time_now = self._time()
tag = datetime.fromtimestamp(int(time_now)).isoformat()
timestamp = int(time_now * 1e3)
request.payload = {
"op": "login",
"tag": tag,
"data": {
"apiKey": self.api_key,
"timestamp": timestamp,
"signature": self._generate_signature_ws(timestamp),
}
}
return request
