def _create_user_stream_data_source(self) -> UserStreamTrackerDataSource:
return BitmartAPIUserStreamDataSource(
auth=self._auth,
trading_pairs=self._trading_pairs,
connector=self,
api_factory=self._web_assistants_factory,
)
def setUp(self) -> None:
super().setUp()
self.listening_task = None
self.log_records = []
throttler = AsyncThrottler(CONSTANTS.RATE_LIMITS)
auth_assistant = BitmartAuth(api_key=self.api_key,
secret_key=self.secret,
memo=self.memo)
self.data_source = BitmartAPIUserStreamDataSource(auth_assistant, throttler)
self.data_source.logger().setLevel(1)
self.data_source.logger().addHandler(self)
self.data_source._trading_pairs = ["HBOT-USDT"]
self.mocking_assistant = NetworkMockingAssistant()
def data_source(self) -> UserStreamTrackerDataSource:
"""
*required
Initializes a user stream data source (user specific order diffs from live socket stream)
:return: OrderBookTrackerDataSource
"""
if not self._data_source:
self._data_source = BitmartAPIUserStreamDataSource(
throttler=self._throttler,
bitmart_auth=self._bitmart_auth,
trading_pairs=self._trading_pairs)
return self._data_source
def __init__(self,
throttler: AsyncThrottler,
bitmart_auth: Optional[BitmartAuth] = None,
trading_pairs: Optional[List[str]] = None,
api_factory: Optional[WebAssistantsFactory] = None):
self._api_factory = api_factory
self._bitmart_auth: BitmartAuth = bitmart_auth
self._trading_pairs: List[str] = trading_pairs or []
self._throttler = throttler
super().__init__(data_source=BitmartAPIUserStreamDataSource(
throttler=self._throttler,
bitmart_auth=self._bitmart_auth,
trading_pairs=self._trading_pairs,
api_factory=self._api_factory))
def setUp(self) -> None:
super().setUp()
self.log_records = []
self.listening_task: Optional[asyncio.Task] = None
self.mocking_assistant = NetworkMockingAssistant()
self.client_config_map = ClientConfigAdapter(ClientConfigMap())
self.time_synchronizer = MagicMock()
self.time_synchronizer.time.return_value = 1640001112.223
self.auth = BitmartAuth(api_key="test_api_key",
secret_key="test_secret_key",
memo="test_memo",
time_provider=self.time_synchronizer)
self.connector = BitmartExchange(
client_config_map=self.client_config_map,
bitmart_api_key="test_api_key",
bitmart_secret_key="test_secret_key",
bitmart_memo="test_memo",
trading_pairs=[self.trading_pair],
trading_required=False,
)
self.connector._web_assistants_factory._auth = self.auth
self.data_source = BitmartAPIUserStreamDataSource(
auth=self.auth,
trading_pairs=[self.trading_pair],
connector=self.connector,
api_factory=self.connector._web_assistants_factory)
self.data_source.logger().setLevel(1)
self.data_source.logger().addHandler(self)
self.connector._set_trading_pair_symbol_map(
bidict({self.ex_trading_pair: self.trading_pair}))
class BitmartAPIUserStreamDataSourceTests(unittest.TestCase):
# the level is required to receive logs from the data source logger
level = 0
@classmethod
def setUpClass(cls) -> None:
super().setUpClass()
cls.ev_loop = asyncio.get_event_loop()
cls.base_asset = "COINALPHA"
cls.quote_asset = "HBOT"
cls.trading_pair = f"{cls.base_asset}-{cls.quote_asset}"
cls.ex_trading_pair = f"{cls.base_asset}_{cls.quote_asset}"
def setUp(self) -> None:
super().setUp()
self.log_records = []
self.listening_task: Optional[asyncio.Task] = None
self.mocking_assistant = NetworkMockingAssistant()
self.client_config_map = ClientConfigAdapter(ClientConfigMap())
self.time_synchronizer = MagicMock()
self.time_synchronizer.time.return_value = 1640001112.223
self.auth = BitmartAuth(api_key="test_api_key",
secret_key="test_secret_key",
memo="test_memo",
time_provider=self.time_synchronizer)
self.connector = BitmartExchange(
client_config_map=self.client_config_map,
bitmart_api_key="test_api_key",
bitmart_secret_key="test_secret_key",
bitmart_memo="test_memo",
trading_pairs=[self.trading_pair],
trading_required=False,
)
self.connector._web_assistants_factory._auth = self.auth
self.data_source = BitmartAPIUserStreamDataSource(
auth=self.auth,
trading_pairs=[self.trading_pair],
connector=self.connector,
api_factory=self.connector._web_assistants_factory)
self.data_source.logger().setLevel(1)
self.data_source.logger().addHandler(self)
self.connector._set_trading_pair_symbol_map(
bidict({self.ex_trading_pair: self.trading_pair}))
def tearDown(self) -> None:
self.listening_task and self.listening_task.cancel()
super().tearDown()
def handle(self, record):
self.log_records.append(record)
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
def _is_logged(self, log_level: str, message: str) -> bool:
return any(
record.levelname == log_level and record.getMessage() == message
for record in self.log_records)
def _raise_exception(self, exception_class):
raise exception_class
@patch("aiohttp.ClientSession.ws_connect", new_callable=AsyncMock)
def test_listen_for_user_stream_subscribes_to_orders_events(
self, ws_connect_mock):
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
successful_login_response = {"event": "login"}
result_subscribe_orders = {
"event": "subscribe",
"topic": CONSTANTS.PRIVATE_ORDER_PROGRESS_CHANNEL_NAME,
}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=ws_connect_mock.return_value,
message=json.dumps(successful_login_response))
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=ws_connect_mock.return_value,
message=json.dumps(result_subscribe_orders))
output_queue = asyncio.Queue()
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(output=output_queue))
self.mocking_assistant.run_until_all_aiohttp_messages_delivered(
ws_connect_mock.return_value)
sent_messages = self.mocking_assistant.json_messages_sent_through_websocket(
websocket_mock=ws_connect_mock.return_value)
self.assertEqual(2, len(sent_messages))
expected_login = {
"op":
"login",
"args": [
"test_api_key",
str(int(self.time_synchronizer.time() * 1e3)),
"f0f176c799346a7730c9c237a09d14742971f3ab59848dde75ef1ac95b04c4e5"
] # noqa: mock
}
self.assertEqual(expected_login, sent_messages[0])
expected_orders_subscription = {
"op":
"subscribe",
"args": [
f"{CONSTANTS.PRIVATE_ORDER_PROGRESS_CHANNEL_NAME}:{self.ex_trading_pair}"
]
}
self.assertEqual(expected_orders_subscription, sent_messages[1])
self.assertTrue(
self._is_logged(
"INFO",
"Subscribed to private account and orders channels..."))
@patch("aiohttp.ClientSession.ws_connect", new_callable=AsyncMock)
def test_listen_for_user_stream_logs_error_when_login_fails(
self, ws_connect_mock):
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
erroneous_login_response = {"event": "login", "errorCode": "4001"}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=ws_connect_mock.return_value,
message=json.dumps(erroneous_login_response))
output_queue = asyncio.Queue()
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(output=output_queue))
self.mocking_assistant.run_until_all_aiohttp_messages_delivered(
ws_connect_mock.return_value)
self.assertTrue(
self._is_logged(
"ERROR",
"Error authenticating the private websocket connection"))
self.assertTrue(
self._is_logged(
"ERROR",
"Unexpected error while listening to user stream. Retrying after 5 seconds..."
))
@patch("aiohttp.ClientSession.ws_connect", new_callable=AsyncMock)
def test_listen_for_user_stream_does_not_queue_invalid_payload(
self, mock_ws):
mock_ws.return_value = self.mocking_assistant.create_websocket_mock()
successful_login_response = {"event": "login"}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=mock_ws.return_value,
message=json.dumps(successful_login_response))
event_without_data = {
"table": CONSTANTS.PRIVATE_ORDER_PROGRESS_CHANNEL_NAME
}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=mock_ws.return_value,
message=json.dumps(event_without_data))
event_without_table = {
"data": [{
"symbol": self.ex_trading_pair,
"side": "buy",
"type": "market",
"notional": "",
"size": "1.0000000000",
"ms_t": "1609926028000",
"price": "46100.0000000000",
"filled_notional": "46100.0000000000",
"filled_size": "1.0000000000",
"margin_trading": "0",
"state": "4",
"order_id": "2147857398",
"order_type": "0",
"last_fill_time": "1609926039226",
"last_fill_price": "46100.00000",
"last_fill_count": "1.00000",
"exec_type": "M",
"detail_id": "256348632",
"client_order_id": "order4872191"
}],
}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=mock_ws.return_value,
message=json.dumps(event_without_table))
msg_queue = asyncio.Queue()
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(msg_queue))
self.mocking_assistant.run_until_all_aiohttp_messages_delivered(
mock_ws.return_value)
self.assertEqual(0, msg_queue.qsize())
@patch("aiohttp.ClientSession.ws_connect", new_callable=AsyncMock)
@patch(
"hummingbot.core.data_type.user_stream_tracker_data_source.UserStreamTrackerDataSource._sleep"
)
def test_listen_for_user_stream_connection_failed(self, sleep_mock,
mock_ws):
mock_ws.side_effect = Exception("TEST ERROR")
sleep_mock.side_effect = asyncio.CancelledError
msg_queue = asyncio.Queue()
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(msg_queue))
try:
self.async_run_with_timeout(self.listening_task)
except asyncio.CancelledError:
pass
self.assertTrue(
self._is_logged(
"ERROR",
"Unexpected error while listening to user stream. Retrying after 5 seconds..."
))
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_canceled_when_cancel_exception_during_initialization(
self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.side_effect = asyncio.CancelledError
with self.assertRaises(asyncio.CancelledError):
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(messages))
self.ev_loop.run_until_complete(self.listening_task)
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_canceled_when_cancel_exception_during_authentication(
self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
ws_connect_mock.return_value.receive.side_effect = asyncio.CancelledError
with self.assertRaises(asyncio.CancelledError):
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(messages))
self.ev_loop.run_until_complete(self.listening_task)
def test_subscribe_channels_raises_cancel_exception(self):
ws_assistant = AsyncMock()
ws_assistant.send.side_effect = asyncio.CancelledError
with self.assertRaises(asyncio.CancelledError):
self.listening_task = self.ev_loop.create_task(
self.data_source._subscribe_channels(ws_assistant))
self.ev_loop.run_until_complete(self.listening_task)
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
@patch(
"hummingbot.core.data_type.user_stream_tracker_data_source.UserStreamTrackerDataSource._sleep"
)
def test_listening_process_logs_exception_during_events_subscription(
self, sleep_mock, mock_ws):
self.connector._set_trading_pair_symbol_map({})
messages = asyncio.Queue()
sleep_mock.side_effect = asyncio.CancelledError
mock_ws.return_value = self.mocking_assistant.create_websocket_mock()
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_aiohttp_message(
mock_ws.return_value, json.dumps({"event": "login"}))
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(messages))
try:
self.async_run_with_timeout(self.listening_task, timeout=3)
except asyncio.CancelledError:
pass
self.assertTrue(
self._is_logged(
"ERROR",
"Unexpected error occurred subscribing to order book trading and delta streams..."
))
self.assertTrue(
self._is_logged(
"ERROR",
"Unexpected error while listening to user stream. Retrying after 5 seconds..."
))
@patch("aiohttp.ClientSession.ws_connect", new_callable=AsyncMock)
def test_listen_for_user_stream_processes_order_event(self, mock_ws):
mock_ws.return_value = self.mocking_assistant.create_websocket_mock()
successful_login_response = {"event": "login"}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=mock_ws.return_value,
message=json.dumps(successful_login_response))
order_event = {
"data": [{
"symbol": self.ex_trading_pair,
"side": "buy",
"type": "market",
"notional": "",
"size": "1.0000000000",
"ms_t": "1609926028000",
"price": "46100.0000000000",
"filled_notional": "46100.0000000000",
"filled_size": "1.0000000000",
"margin_trading": "0",
"state": "4",
"order_id": "2147857398",
"order_type": "0",
"last_fill_time": "1609926039226",
"last_fill_price": "46100.00000",
"last_fill_count": "1.00000",
"exec_type": "M",
"detail_id": "256348632",
"client_order_id": "order4872191"
}],
"table":
CONSTANTS.PRIVATE_ORDER_PROGRESS_CHANNEL_NAME
}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=mock_ws.return_value,
message=json.dumps(order_event))
msg_queue = asyncio.Queue()
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(msg_queue))
self.mocking_assistant.run_until_all_aiohttp_messages_delivered(
mock_ws.return_value)
self.assertEqual(1, msg_queue.qsize())
order_event_message = msg_queue.get_nowait()
self.assertEqual(order_event, order_event_message)
@patch("aiohttp.ClientSession.ws_connect", new_callable=AsyncMock)
def test_listen_for_user_stream_processes_compressed_order_event(
self, mock_ws):
mock_ws.return_value = self.mocking_assistant.create_websocket_mock()
successful_login_response = {"event": "login"}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=mock_ws.return_value,
message=json.dumps(successful_login_response))
order_event = {
"data": [{
"symbol": self.ex_trading_pair,
"side": "buy",
"type": "market",
"notional": "",
"size": "1.0000000000",
"ms_t": "1609926028000",
"price": "46100.0000000000",
"filled_notional": "46100.0000000000",
"filled_size": "1.0000000000",
"margin_trading": "0",
"state": "4",
"order_id": "2147857398",
"order_type": "0",
"last_fill_time": "1609926039226",
"last_fill_price": "46100.00000",
"last_fill_count": "1.00000",
"exec_type": "M",
"detail_id": "256348632",
"client_order_id": "order4872191"
}],
"table":
CONSTANTS.PRIVATE_ORDER_PROGRESS_CHANNEL_NAME
}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=mock_ws.return_value,
message=bitmart_utils.compress_ws_message(json.dumps(order_event)),
message_type=WSMsgType.BINARY)
msg_queue = asyncio.Queue()
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(msg_queue))
self.mocking_assistant.run_until_all_aiohttp_messages_delivered(
mock_ws.return_value)
self.assertEqual(1, msg_queue.qsize())
order_event_message = msg_queue.get_nowait()
self.assertEqual(order_event, order_event_message)
@patch("aiohttp.ClientSession.ws_connect", new_callable=AsyncMock)
def test_listen_for_user_stream_logs_details_for_order_event_with_errors(
self, mock_ws):
mock_ws.return_value = self.mocking_assistant.create_websocket_mock()
successful_login_response = {"event": "login"}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=mock_ws.return_value,
message=json.dumps(successful_login_response))
order_event = {
"errorCode":
"4001",
"errorMessage":
"Error",
"data": [{
"symbol": self.ex_trading_pair,
"side": "buy",
"type": "market",
"notional": "",
"size": "1.0000000000",
"ms_t": "1609926028000",
"price": "46100.0000000000",
"filled_notional": "46100.0000000000",
"filled_size": "1.0000000000",
"margin_trading": "0",
"state": "4",
"order_id": "2147857398",
"order_type": "0",
"last_fill_time": "1609926039226",
"last_fill_price": "46100.00000",
"last_fill_count": "1.00000",
"exec_type": "M",
"detail_id": "256348632",
"client_order_id": "order4872191"
}],
"table":
CONSTANTS.PRIVATE_ORDER_PROGRESS_CHANNEL_NAME
}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=mock_ws.return_value,
message=json.dumps(order_event))
msg_queue = asyncio.Queue()
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(msg_queue))
self.mocking_assistant.run_until_all_aiohttp_messages_delivered(
mock_ws.return_value)
self.assertEqual(0, msg_queue.qsize())
self.assertTrue(
self._is_logged(
"ERROR",
"Unexpected error while listening to user stream. Retrying after 5 seconds..."
))
@patch("aiohttp.ClientSession.ws_connect", new_callable=AsyncMock)
def test_listen_for_user_stream_logs_details_for_invalid_event_message(
self, mock_ws):
mock_ws.return_value = self.mocking_assistant.create_websocket_mock()
successful_login_response = {"event": "login"}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=mock_ws.return_value,
message=json.dumps(successful_login_response))
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=mock_ws.return_value,
message="invalid message content")
msg_queue = asyncio.Queue()
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(msg_queue))
self.mocking_assistant.run_until_all_aiohttp_messages_delivered(
mock_ws.return_value)
self.assertEqual(0, msg_queue.qsize())
self.assertTrue(
self._is_logged(
"WARNING",
"Invalid event message received through the order book data source connection (invalid message content)"
))
class BitmartAPIUserStreamDataSourceTests(unittest.TestCase):
# the level is required to receive logs from the data source logger
level = 0
@classmethod
def setUpClass(cls) -> None:
super().setUpClass()
cls.api_key = 'testAPIKey'
cls.secret = 'testSecret'
cls.memo = '001'
cls.account_id = 528
cls.username = '******'
cls.oms_id = 1
cls.ev_loop = asyncio.get_event_loop()
def setUp(self) -> None:
super().setUp()
self.listening_task = None
self.log_records = []
throttler = AsyncThrottler(CONSTANTS.RATE_LIMITS)
auth_assistant = BitmartAuth(api_key=self.api_key,
secret_key=self.secret,
memo=self.memo)
self.data_source = BitmartAPIUserStreamDataSource(
auth_assistant, throttler)
self.data_source.logger().setLevel(1)
self.data_source.logger().addHandler(self)
self.data_source._trading_pairs = ["HBOT-USDT"]
self.mocking_assistant = NetworkMockingAssistant()
def tearDown(self) -> None:
self.listening_task and self.listening_task.cancel()
super().tearDown()
def handle(self, record):
self.log_records.append(record)
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
def _is_logged(self, log_level: str, message: str) -> bool:
return any(
record.levelname == log_level and record.getMessage() == message
for record in self.log_records)
def _raise_exception(self, exception_class):
raise exception_class
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_authenticates_and_subscribes_to_events(
self, ws_connect_mock):
mock_response: Dict[Any] = {
"data": [{
"symbol": "BTC_USDT",
"side": "buy",
"type": "market",
"notional": "",
"size": "1.0000000000",
"ms_t": "1609926028000",
"price": "46100.0000000000",
"filled_notional": "46100.0000000000",
"filled_size": "1.0000000000",
"margin_trading": "0",
"state": "2",
"order_id": "2147857398",
"order_type": "0",
"last_fill_time": "1609926039226",
"last_fill_price": "46100.00000",
"last_fill_count": "1.00000"
}],
"table":
"spot/user/order"
}
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
initial_last_recv_time = self.data_source.last_recv_time
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop, messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_aiohttp_message(
ws_connect_mock.return_value, json.dumps({"event": "login"}))
# Add a dummy message for the websocket to read and include in the "messages" queue
self.mocking_assistant.add_websocket_aiohttp_message(
ws_connect_mock.return_value, json.dumps(mock_response))
first_received_message = self.ev_loop.run_until_complete(
messages.get())
self.assertEqual(mock_response, first_received_message)
self.assertTrue(
self._is_logged('INFO', "Authenticating to User Stream..."))
self.assertTrue(
self._is_logged('INFO',
"Successfully authenticated to User Stream."))
self.assertTrue(
self._is_logged(
'INFO', "Successfully subscribed to all Private channels."))
sent_messages = self.mocking_assistant.json_messages_sent_through_websocket(
ws_connect_mock.return_value)
self.assertEqual(2, len(sent_messages))
auth_req = sent_messages[0]
sub_req = sent_messages[1]
self.assertTrue("op" in auth_req and "args" in auth_req
and "testAPIKey" in auth_req["args"])
self.assertEqual(
{
"op": "subscribe",
"args": ["spot/user/order:HBOT_USDT"]
}, sub_req)
self.assertGreater(self.data_source.last_recv_time,
initial_last_recv_time)
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_fails_when_authentication_fails(
self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
# Make the close function raise an exception to finish the execution
ws_connect_mock.return_value.closed = False
ws_connect_mock.return_value.close.side_effect = lambda: self._raise_exception(
Exception)
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop, messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_aiohttp_message(
ws_connect_mock.return_value,
json.dumps({
"errorCode": "test code",
"errorMessage": "test err message"
}))
try:
self.ev_loop.run_until_complete(self.listening_task)
except Exception:
pass
self.assertTrue(
self._is_logged(
"ERROR",
"WebSocket login errored with message: test err message"))
self.assertTrue(
self._is_logged(
"ERROR", "Error occurred when authenticating to user stream."))
self.assertTrue(
self._is_logged(
"ERROR", "Unexpected error with BitMart WebSocket connection. "
"Retrying after 30 seconds..."))
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_canceled_when_cancel_exception_during_initialization(
self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.side_effect = asyncio.CancelledError
with self.assertRaises(asyncio.CancelledError):
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(
self.ev_loop, messages))
self.ev_loop.run_until_complete(self.listening_task)
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_canceled_when_cancel_exception_during_authentication(
self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
ws_connect_mock.return_value.send_json.side_effect = lambda sent_message: (
self._raise_exception(asyncio.CancelledError)
if "testAPIKey" in sent_message["args"] else self.mocking_assistant
._sent_websocket_json_messages[ws_connect_mock.return_value
].append(sent_message))
with self.assertRaises(asyncio.CancelledError):
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(
self.ev_loop, messages))
self.ev_loop.run_until_complete(self.listening_task)
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_canceled_when_cancel_exception_during_events_subscription(
self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
ws_connect_mock.return_value.send_json.side_effect = lambda sent_message: (
self._raise_exception(asyncio.CancelledError)
if "spot/user/order:HBOT_USDT" in sent_message["args"] else self.
mocking_assistant._sent_websocket_json_messages[
ws_connect_mock.return_value].append(sent_message))
with self.assertRaises(asyncio.CancelledError):
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(
self.ev_loop, messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_aiohttp_message(
ws_connect_mock.return_value, json.dumps({"event": "login"}))
self.ev_loop.run_until_complete(self.listening_task)
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_logs_exception_details_during_initialization(
self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.side_effect = Exception
with self.assertRaises(Exception):
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(
self.ev_loop, messages))
try:
self.async_run_with_timeout(self.listening_task)
except asyncio.TimeoutError:
raise
self.assertTrue(
self._is_logged(
"ERROR",
"Unexpected error with BitMart WebSocket connection. Retrying after 30 seconds..."
))
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_logs_exception_details_during_authentication(
self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
ws_connect_mock.return_value.send_json.side_effect = lambda sent_message: (
self._raise_exception(Exception)
if "testAPIKey" in sent_message["args"] else self.mocking_assistant
._sent_websocket_json_messages[ws_connect_mock.return_value
].append(sent_message))
# Make the close function raise an exception to finish the execution
ws_connect_mock.return_value.closed = False
ws_connect_mock.return_value.close.side_effect = lambda: self._raise_exception(
Exception)
try:
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(
self.ev_loop, messages))
self.ev_loop.run_until_complete(self.listening_task)
except Exception:
pass
self.assertTrue(
self._is_logged(
"ERROR", "Error occurred when authenticating to user stream."))
self.assertTrue(
self._is_logged(
"ERROR", "Unexpected error with BitMart WebSocket connection. "
"Retrying after 30 seconds..."))
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_logs_exception_during_events_subscription(
self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
ws_connect_mock.return_value.send_json.side_effect = lambda sent_message: (
self._raise_exception(Exception)
if "spot/user/order:HBOT_USDT" in sent_message["args"] else self.
mocking_assistant._sent_websocket_json_messages[
ws_connect_mock.return_value].append(sent_message))
# Make the close function raise an exception to finish the execution
ws_connect_mock.return_value.closed = False
ws_connect_mock.return_value.close.side_effect = lambda: self._raise_exception(
Exception)
try:
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(
self.ev_loop, messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_aiohttp_message(
ws_connect_mock.return_value, json.dumps({"event": "login"}))
self.async_run_with_timeout(self.listening_task)
except Exception:
pass
self.assertTrue(
self._is_logged(
"ERROR",
"Error occured during subscribing to Bitmart private channels."
))
self.assertTrue(
self._is_logged(
"ERROR", "Unexpected error with BitMart WebSocket connection. "
"Retrying after 30 seconds..."))
class BitmartAPIUserStreamDataSourceTests(unittest.TestCase):
# the level is required to receive logs from the data source logger
level = 0
@classmethod
def setUpClass(cls) -> None:
super().setUpClass()
cls.api_key = 'testAPIKey'
cls.secret = 'testSecret'
cls.memo = '001'
cls.account_id = 528
cls.username = '******'
cls.oms_id = 1
cls.ev_loop = asyncio.get_event_loop()
def setUp(self) -> None:
super().setUp()
self.listening_task = None
self.log_records = []
throttler = AsyncThrottler(CONSTANTS.RATE_LIMITS)
auth_assistant = BitmartAuth(api_key=self.api_key,
secret_key=self.secret,
memo=self.memo)
self.data_source = BitmartAPIUserStreamDataSource(auth_assistant, throttler)
self.data_source.logger().setLevel(1)
self.data_source.logger().addHandler(self)
self.data_source._trading_pairs = ["HBOT-USDT"]
self.mocking_assistant = NetworkMockingAssistant()
def tearDown(self) -> None:
self.listening_task and self.listening_task.cancel()
super().tearDown()
def handle(self, record):
self.log_records.append(record)
def _is_logged(self, log_level: str, message: str) -> bool:
return any(record.levelname == log_level and record.getMessage() == message
for record in self.log_records)
def _raise_exception(self, exception_class):
raise exception_class
@patch('websockets.connect', new_callable=AsyncMock)
def test_listening_process_authenticates_and_subscribes_to_events(self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock()
initial_last_recv_time = self.data_source.last_recv_time
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop,
messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_text_message(
ws_connect_mock.return_value,
json.dumps({"event": "login"}))
# Add a dummy message for the websocket to read and include in the "messages" queue
self.mocking_assistant.add_websocket_text_message(ws_connect_mock.return_value, json.dumps('dummyMessage'))
first_received_message = self.ev_loop.run_until_complete(messages.get())
self.assertEqual('dummyMessage', first_received_message)
self.assertTrue(self._is_logged('INFO', "Authenticating to User Stream..."))
self.assertTrue(self._is_logged('INFO', "Successfully authenticated to User Stream."))
self.assertTrue(self._is_logged('INFO', "Successfully subscribed to all Private channels."))
sent_messages = self.mocking_assistant.text_messages_sent_through_websocket(ws_connect_mock.return_value)
self.assertEqual(2, len(sent_messages))
auth_req = json.loads(sent_messages[0])
sub_req = json.loads(sent_messages[1])
self.assertTrue("op" in auth_req and "args" in auth_req and "testAPIKey" in auth_req["args"])
self.assertEqual({"op": "subscribe", "args": ["spot/user/order:HBOT_USDT"]},
sub_req)
self.assertGreater(self.data_source.last_recv_time, initial_last_recv_time)
@patch('websockets.connect', new_callable=AsyncMock)
def test_listening_process_fails_when_authentication_fails(self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock()
# Make the close function raise an exception to finish the execution
ws_connect_mock.return_value.close.side_effect = lambda: self._raise_exception(Exception)
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop,
messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_text_message(
ws_connect_mock.return_value,
json.dumps({"errorCode": "test code", "errorMessage": "test err message"})
)
try:
self.ev_loop.run_until_complete(self.listening_task)
except Exception:
pass
self.assertTrue(self._is_logged("ERROR", "WebSocket login errored with message: test err message"))
self.assertTrue(self._is_logged("ERROR", "Error occurred when authenticating to user stream."))
self.assertTrue(self._is_logged("ERROR", "Unexpected error with BitMart WebSocket connection. "
"Retrying after 30 seconds..."))
@patch('websockets.connect', new_callable=AsyncMock)
def test_listening_process_canceled_when_cancel_exception_during_initialization(self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.side_effect = asyncio.CancelledError
with self.assertRaises(asyncio.CancelledError):
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop,
messages))
self.ev_loop.run_until_complete(self.listening_task)
@patch('websockets.connect', new_callable=AsyncMock)
def test_listening_process_canceled_when_cancel_exception_during_authentication(self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock()
ws_connect_mock.return_value.send.side_effect = lambda sent_message: (
self._raise_exception(asyncio.CancelledError)
if "testAPIKey" in sent_message
else self.mocking_assistant._sent_websocket_text_messages[ws_connect_mock.return_value].append(
sent_message))
with self.assertRaises(asyncio.CancelledError):
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop,
messages))
self.ev_loop.run_until_complete(self.listening_task)
@patch('websockets.connect', new_callable=AsyncMock)
def test_listening_process_canceled_when_cancel_exception_during_events_subscription(self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock()
ws_connect_mock.return_value.send.side_effect = lambda sent_message: (
self._raise_exception(asyncio.CancelledError)
if "order:HBOT_USDT" in sent_message
else self.mocking_assistant._sent_websocket_text_messages[ws_connect_mock.return_value].append(sent_message)
)
with self.assertRaises(asyncio.CancelledError):
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop,
messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_text_message(
ws_connect_mock.return_value,
json.dumps({"event": "login"})
)
self.ev_loop.run_until_complete(self.listening_task)
@patch('websockets.connect', new_callable=AsyncMock)
def test_listening_process_logs_exception_details_during_initialization(self, ws_connect_mock):
ws_connect_mock.side_effect = Exception
with self.assertRaises(Exception):
self.listening_task = self.ev_loop.create_task(self.data_source._init_websocket_connection())
self.ev_loop.run_until_complete(self.listening_task)
self.assertTrue(self._is_logged("NETWORK", "Unexpected error occured with BitMart WebSocket Connection"))
@patch('websockets.connect', new_callable=AsyncMock)
def test_listening_process_logs_exception_details_during_authentication(self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock()
ws_connect_mock.return_value.send.side_effect = lambda sent_message: (
self._raise_exception(Exception)
if "testAPIKey" in sent_message
else self.mocking_assistant._sent_websocket_text_messages[ws_connect_mock.return_value].append(sent_message))
# Make the close function raise an exception to finish the execution
ws_connect_mock.return_value.close.side_effect = lambda: self._raise_exception(Exception)
try:
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop,
messages))
self.ev_loop.run_until_complete(self.listening_task)
except Exception:
pass
self.assertTrue(self._is_logged("ERROR", "Error occurred when authenticating to user stream."))
self.assertTrue(self._is_logged("ERROR", "Unexpected error with BitMart WebSocket connection. "
"Retrying after 30 seconds..."))
@patch('websockets.connect', new_callable=AsyncMock)
def test_listening_process_logs_exception_during_events_subscription(self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock()
ws_connect_mock.return_value.send.side_effect = lambda sent_message: (
self._raise_exception(Exception)
if "order:HBOT_USDT" in sent_message
else self.mocking_assistant._sent_websocket_text_messages[ws_connect_mock.return_value].append(sent_message)
)
# Make the close function raise an exception to finish the execution
ws_connect_mock.return_value.close.side_effect = lambda: self._raise_exception(Exception)
try:
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop,
messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_text_message(
ws_connect_mock.return_value,
json.dumps({"event": "login"}))
self.ev_loop.run_until_complete(self.listening_task)
except Exception:
pass
self.assertTrue(self._is_logged("ERROR", "Error occured during subscribing to Bitmart private channels."))
self.assertTrue(self._is_logged("ERROR", "Unexpected error with BitMart WebSocket connection. "
"Retrying after 30 seconds..."))
@patch('websockets.connect', new_callable=AsyncMock)
def test_listening_process_invalid_json_message_logged(self, ws_connect_mock):
messages = asyncio.Queue()
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock()
# Make the close function raise an exception to finish the execution
ws_connect_mock.return_value.close.side_effect = lambda: self._raise_exception(Exception)
try:
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop,
messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_text_message(
ws_connect_mock.return_value,
json.dumps({"event": "login"}))
# Add invalid json message
self.mocking_assistant.add_websocket_text_message(ws_connect_mock.return_value, 'invalid message')
self.ev_loop.run_until_complete(self.listening_task)
except Exception:
pass
self.assertTrue(self._is_logged("ERROR", "Unexpected error when parsing BitMart user_stream message. "))
self.assertTrue(self._is_logged("ERROR", "Unexpected error with BitMart WebSocket connection. "
"Retrying after 30 seconds..."))
@patch('websockets.connect', new_callable=AsyncMock)
def test_listen_for_user_stream_inner_messages_recv_timeout(self, ws_connect_mock):
self.data_source.MESSAGE_TIMEOUT = 0.1
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock()
ws_connect_mock.return_value.ping.side_effect = lambda: done_callback_event.set()
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_text_message(
ws_connect_mock.return_value,
json.dumps({"event": "login"}))
done_callback_event = asyncio.Event()
message_queue = asyncio.Queue()
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop,
message_queue)
)
self.ev_loop.run_until_complete(done_callback_event.wait())
@patch('websockets.connect', new_callable=AsyncMock)
def test_listen_for_user_stream_inner_messages_recv_timeout_ping_timeout(self, ws_connect_mock):
self.data_source.PING_TIMEOUT = 0.1
self.data_source.MESSAGE_TIMEOUT = 0.1
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock()
ws_connect_mock.return_value.close.side_effect = lambda: done_callback_event.set()
ws_connect_mock.return_value.ping.side_effect = NetworkMockingAssistant.async_partial(
self.mocking_assistant._get_next_websocket_text_message, ws_connect_mock.return_value
)
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_text_message(
ws_connect_mock.return_value,
json.dumps({"event": "login"}))
done_callback_event = asyncio.Event()
message_queue = asyncio.Queue()
self.listening_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop,
message_queue)
)
self.ev_loop.run_until_complete(done_callback_event.wait())
self.assertTrue(self._is_logged("WARNING", "WebSocket ping timed out. Going to reconnect..."))