class ProbitAPIUserStreamDataSourceTest(unittest.TestCase):
# logging.Level required to receive logs from the data source logger
level = 0
@classmethod
def setUpClass(cls) -> None:
cls.base_asset = "BTC"
cls.quote_asset = "USDT"
cls.trading_pair = f"{cls.base_asset}-{cls.quote_asset}"
def setUp(self) -> None:
super().setUp()
self.ev_loop = asyncio.get_event_loop()
self.api_key = "someKey"
self.api_secret = "someSecret"
self.auth = ProbitAuth(self.api_key, self.api_secret)
self.data_source = ProbitAPIUserStreamDataSource(
self.auth, trading_pairs=[self.trading_pair])
self.data_source.logger().setLevel(1)
self.data_source.logger().addHandler(self)
self.log_records = []
self.mocking_assistant = NetworkMockingAssistant()
self.async_task: Optional[asyncio.Task] = None
def tearDown(self) -> None:
self.async_task and self.async_task.cancel()
super().tearDown()
def handle(self, record):
self.log_records.append(record)
def check_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 async_run_with_timeout(self, coroutine: Awaitable, timeout: float = 1):
ret = self.ev_loop.run_until_complete(
asyncio.wait_for(coroutine, timeout))
return ret
@patch("aiohttp.client.ClientSession.ws_connect", new_callable=AsyncMock)
@patch(
"hummingbot.connector.exchange.probit.probit_auth.ProbitAuth.get_ws_auth_payload",
new_callable=AsyncMock,
)
def test_listen_for_user_stream(self, get_ws_auth_payload_mock,
ws_connect_mock):
auth_msg = {"type": "authorization", "token": "someToken"}
get_ws_auth_payload_mock.return_value = auth_msg
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value,
message={"result": "ok"} # authentication
)
self.mocking_assistant.add_websocket_aiohttp_message(
ws_connect_mock.return_value,
message=json.dumps({"my_msg": "test"}) # first message
)
output_queue = asyncio.Queue()
self.async_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(output_queue))
self.mocking_assistant.run_until_all_json_messages_delivered(
ws_connect_mock.return_value)
self.mocking_assistant.run_until_all_aiohttp_messages_delivered(
ws_connect_mock.return_value)
self.assertFalse(output_queue.empty())
sent_text_msgs = self.mocking_assistant.text_messages_sent_through_websocket(
ws_connect_mock.return_value)
self.assertEqual(auth_msg, json.loads(sent_text_msgs[0]))
sent_json_msgs = self.mocking_assistant.json_messages_sent_through_websocket(
ws_connect_mock.return_value)
for sent_json_msg in sent_json_msgs:
self.assertEqual("subscribe", sent_json_msg["type"])
self.assertIn(sent_json_msg["channel"],
CONSTANTS.WS_PRIVATE_CHANNELS)
CONSTANTS.WS_PRIVATE_CHANNELS.remove(sent_json_msg["channel"])
self.assertEqual(0, len(CONSTANTS.WS_PRIVATE_CHANNELS))
self.assertNotEqual(0, self.data_source.last_recv_time)
@patch("aiohttp.client.ClientSession.ws_connect")
@patch(
"hummingbot.connector.exchange.probit.probit_api_user_stream_data_source.ProbitAPIUserStreamDataSource._sleep",
new_callable=AsyncMock,
)
def test_listen_for_user_stream_attempts_again_on_exception(
self, sleep_mock, ws_connect_mock):
called_event = asyncio.Event()
async def _sleep(delay):
called_event.set()
await asyncio.sleep(delay)
sleep_mock.side_effect = _sleep
ws_connect_mock.side_effect = Exception
self.async_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(asyncio.Queue()))
self.async_run_with_timeout(called_event.wait())
self.check_is_logged(
log_level="ERROR",
message=
"Unexpected error with Probit WebSocket connection. Retrying after 30 seconds...",
)
@patch("aiohttp.client.ClientSession.ws_connect")
def test_listen_for_user_stream_stops_on_asyncio_cancelled_error(
self, ws_connect_mock):
ws_connect_mock.side_effect = asyncio.CancelledError
with self.assertRaises(asyncio.CancelledError):
self.async_run_with_timeout(
self.data_source.listen_for_user_stream(asyncio.Queue()))
@patch("aiohttp.client.ClientSession.ws_connect", new_callable=AsyncMock)
@patch(
"hummingbot.connector.exchange.probit.probit_auth.ProbitAuth.get_ws_auth_payload",
new_callable=AsyncMock,
)
def test_listen_for_user_stream_registers_ping_msg(
self, get_ws_auth_payload_mock, ws_connect_mock):
auth_msg = {"type": "authorization", "token": "someToken"}
get_ws_auth_payload_mock.return_value = auth_msg
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value,
message={"result": "ok"} # authentication
)
self.mocking_assistant.add_websocket_aiohttp_message(
ws_connect_mock.return_value,
message="",
message_type=WSMsgType.PING)
output_queue = asyncio.Queue()
self.async_task = self.ev_loop.create_task(
self.data_source.listen_for_user_stream(output_queue))
self.mocking_assistant.run_until_all_aiohttp_messages_delivered(
ws_connect_mock.return_value)
self.assertTrue(output_queue.empty())
ws_connect_mock.return_value.pong.assert_called()
class BybitPerpetualUserStreamDataSourceTests(TestCase):
# the level is required to receive logs from the data source loger
level = 0
def setUp(self) -> None:
super().setUp()
self.api_key = 'testAPIKey'
self.secret = 'testSecret'
self.log_records = []
self.listening_task = None
self.data_source = BybitPerpetualUserStreamDataSource(
auth_assistant=BybitPerpetualAuth(api_key=self.api_key,
secret_key=self.secret))
self.data_source.logger().setLevel(1)
self.data_source.logger().addHandler(self)
self.mocking_assistant = NetworkMockingAssistant()
def tearDown(self) -> None:
self.listening_task and self.listening_task.cancel()
if self.data_source._session is not None:
asyncio.get_event_loop().run_until_complete(
self.data_source._session.close())
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 _authentication_response(self, authenticated: bool) -> str:
request = {
"op": "auth",
"args": ['testAPIKey', 'testExpires', 'testSignature']
}
message = {
"success": authenticated,
"ret_msg": "",
"conn_id": "testConnectionID",
"request": request
}
return message
def _subscription_response(self, subscribed: bool,
subscription: str) -> str:
request = {"op": "subscribe", "args": [subscription]}
message = {
"success": subscribed,
"ret_msg": "",
"conn_id": "testConnectionID",
"request": request
}
return message
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):
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 = asyncio.get_event_loop().create_task(
self.data_source._listen_for_user_stream_on_url(
"test_url", messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value, self._authentication_response(True))
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value,
self._subscription_response(
True, CONSTANTS.WS_SUBSCRIPTION_POSITIONS_ENDPOINT_NAME))
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value,
self._subscription_response(
True, CONSTANTS.WS_SUBSCRIPTION_ORDERS_ENDPOINT_NAME))
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value,
self._subscription_response(
True, CONSTANTS.WS_SUBSCRIPTION_EXECUTIONS_ENDPOINT_NAME))
# Add a dummy message for the websocket to read and include in the "messages" queue
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value, json.dumps('dummyMessage'))
asyncio.get_event_loop().run_until_complete(messages.get())
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',
"Successful subscription to the topic ['position'] on test_url"
))
self.assertTrue(
self._is_logged(
"INFO",
"Successful subscription to the topic ['order'] on test_url"))
self.assertTrue(
self._is_logged(
"INFO",
"Successful subscription to the topic ['execution'] on test_url"
))
sent_messages = self.mocking_assistant.json_messages_sent_through_websocket(
ws_connect_mock.return_value)
self.assertEqual(4, len(sent_messages))
authentication_request = sent_messages[0]
subscription_positions_request = sent_messages[1]
subscription_orders_request = sent_messages[2]
subscription_executions_request = sent_messages[3]
self.assertEqual(
CONSTANTS.WS_AUTHENTICATE_USER_ENDPOINT_NAME,
BybitPerpetualWebSocketAdaptor.endpoint_from_message(
authentication_request))
self.assertEqual(
CONSTANTS.WS_SUBSCRIPTION_POSITIONS_ENDPOINT_NAME,
BybitPerpetualWebSocketAdaptor.endpoint_from_message(
subscription_positions_request))
self.assertEqual(
CONSTANTS.WS_SUBSCRIPTION_ORDERS_ENDPOINT_NAME,
BybitPerpetualWebSocketAdaptor.endpoint_from_message(
subscription_orders_request))
self.assertEqual(
CONSTANTS.WS_SUBSCRIPTION_EXECUTIONS_ENDPOINT_NAME,
BybitPerpetualWebSocketAdaptor.endpoint_from_message(
subscription_executions_request))
subscription_positions_payload = BybitPerpetualWebSocketAdaptor.payload_from_message(
subscription_positions_request)
expected_payload = {"op": "subscribe", "args": ["position"]}
self.assertEqual(expected_payload, subscription_positions_payload)
subscription_orders_payload = BybitPerpetualWebSocketAdaptor.payload_from_message(
subscription_orders_request)
expected_payload = {"op": "subscribe", "args": ["order"]}
self.assertEqual(expected_payload, subscription_orders_payload)
subscription_executions_payload = BybitPerpetualWebSocketAdaptor.payload_from_message(
subscription_executions_request)
expected_payload = {"op": "subscribe", "args": ["execution"]}
self.assertEqual(expected_payload, subscription_executions_payload)
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.close.side_effect = lambda: self._raise_exception(
Exception)
self.listening_task = asyncio.get_event_loop().create_task(
self.data_source.listen_for_user_stream(messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value, self._authentication_response(False))
try:
asyncio.get_event_loop().run_until_complete(self.listening_task)
except Exception:
pass
self.assertTrue(
self._is_logged(
"ERROR", "Error occurred when authenticating to user stream "
"(Could not authenticate websocket connection with Bybit Perpetual)"
))
self.assertTrue(
self._is_logged(
"ERROR",
"Unexpected error with Bybit Perpetual WebSocket connection on"
" wss://stream.bybit.com/realtime_private. Retrying in 30 seconds. "
"(Could not authenticate websocket connection with Bybit Perpetual)"
))
@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 = asyncio.get_event_loop().create_task(
self.data_source.listen_for_user_stream(messages))
asyncio.get_event_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 CONSTANTS.WS_AUTHENTICATE_USER_ENDPOINT_NAME in sent_message[
"op"] else self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value, sent_message))
with self.assertRaises(asyncio.CancelledError):
self.listening_task = asyncio.get_event_loop().create_task(
self.data_source.listen_for_user_stream(messages))
asyncio.get_event_loop().run_until_complete(self.listening_task)
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_canceled_when_cancel_exception_during_positions_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 CONSTANTS.
WS_SUBSCRIPTION_POSITIONS_ENDPOINT_NAME in sent_message[
"args"] else self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value, sent_message))
with self.assertRaises(asyncio.CancelledError):
self.listening_task = asyncio.get_event_loop().create_task(
self.data_source.listen_for_user_stream(messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value,
self._authentication_response(True))
asyncio.get_event_loop().run_until_complete(self.listening_task)
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_canceled_when_cancel_exception_during_orders_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 CONSTANTS.WS_SUBSCRIPTION_ORDERS_ENDPOINT_NAME in sent_message[
"args"] else self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value, sent_message))
with self.assertRaises(asyncio.CancelledError):
self.listening_task = asyncio.get_event_loop().create_task(
self.data_source.listen_for_user_stream(messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value,
self._authentication_response(True))
asyncio.get_event_loop().run_until_complete(self.listening_task)
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_canceled_when_cancel_exception_during_executions_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 CONSTANTS.
WS_SUBSCRIPTION_EXECUTIONS_ENDPOINT_NAME in sent_message[
"args"] else self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value, sent_message))
with self.assertRaises(asyncio.CancelledError):
self.listening_task = asyncio.get_event_loop().create_task(
self.data_source.listen_for_user_stream(messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value,
self._authentication_response(True))
asyncio.get_event_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):
ws_connect_mock.side_effect = Exception
with self.assertRaises(Exception):
self.listening_task = asyncio.get_event_loop().create_task(
self.data_source._create_websocket_connection("test_url"))
asyncio.get_event_loop().run_until_complete(self.listening_task)
self.assertTrue(
self._is_logged(
"NETWORK",
"Unexpected error occurred during bybit_perpetual WebSocket Connection on test_url ()"
))
@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 CONSTANTS.WS_AUTHENTICATE_USER_ENDPOINT_NAME in sent_message[
"op"] else self.mocking_assistant.add_websocket_json_message(
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 = asyncio.get_event_loop().create_task(
self.data_source.listen_for_user_stream(messages))
asyncio.get_event_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 Bybit Perpetual WebSocket connection on"
" wss://stream.bybit.com/realtime_private. Retrying in 30 seconds. ()"
))
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_logs_exception_during_positions_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 CONSTANTS.
WS_SUBSCRIPTION_POSITIONS_ENDPOINT_NAME in sent_message[
"args"] else self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value, 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 = asyncio.get_event_loop().create_task(
self.data_source.listen_for_user_stream(messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value,
self._authentication_response(True))
asyncio.get_event_loop().run_until_complete(self.listening_task)
except Exception:
pass
self.assertTrue(
self._is_logged(
"ERROR",
"Error occurred subscribing to bybit_perpetual private channels ()"
))
self.assertTrue(
self._is_logged(
"ERROR",
"Unexpected error with Bybit Perpetual WebSocket connection on"
" wss://stream.bybit.com/realtime_private. Retrying in 30 seconds. ()"
))
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_logs_exception_during_orders_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 CONSTANTS.WS_SUBSCRIPTION_ORDERS_ENDPOINT_NAME in sent_message[
"args"] else self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value, 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 = asyncio.get_event_loop().create_task(
self.data_source.listen_for_user_stream(messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value,
self._authentication_response(True))
asyncio.get_event_loop().run_until_complete(self.listening_task)
except Exception:
pass
self.assertTrue(
self._is_logged(
"ERROR",
"Error occurred subscribing to bybit_perpetual private channels ()"
))
self.assertTrue(
self._is_logged(
"ERROR",
"Unexpected error with Bybit Perpetual WebSocket connection on"
" wss://stream.bybit.com/realtime_private. Retrying in 30 seconds. ()"
))
@patch('aiohttp.ClientSession.ws_connect', new_callable=AsyncMock)
def test_listening_process_logs_exception_during_executions_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 CONSTANTS.
WS_SUBSCRIPTION_EXECUTIONS_ENDPOINT_NAME in sent_message[
"args"] else self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value, 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 = asyncio.get_event_loop().create_task(
self.data_source.listen_for_user_stream(messages))
# Add the authentication response for the websocket
self.mocking_assistant.add_websocket_json_message(
ws_connect_mock.return_value,
self._authentication_response(True))
asyncio.get_event_loop().run_until_complete(self.listening_task)
except Exception:
pass
self.assertTrue(
self._is_logged(
"ERROR",
"Error occurred subscribing to bybit_perpetual private channels ()"
))
self.assertTrue(
self._is_logged(
"ERROR",
"Unexpected error with Bybit Perpetual WebSocket connection on"
" wss://stream.bybit.com/realtime_private. Retrying in 30 seconds. ()"
))
