def setUp(self) -> None:
super().setUp()
self.log_records = []
self.listening_task: Optional[asyncio.Task] = None
self.mocking_assistant = NetworkMockingAssistant()
self.throttler = AsyncThrottler(CONSTANTS.RATE_LIMITS)
self.mock_time_provider = MagicMock()
self.mock_time_provider.time.return_value = 1000
self.auth = KucoinAuth(
self.api_key,
self.api_passphrase,
self.api_secret_key,
time_provider=self.mock_time_provider)
client_config_map = ClientConfigAdapter(ClientConfigMap())
self.connector = KucoinExchange(
client_config_map=client_config_map,
kucoin_api_key="",
kucoin_passphrase="",
kucoin_secret_key="",
trading_pairs=[],
trading_required=False)
self.data_source = KucoinAPIUserStreamDataSource(
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)
def setUp(self) -> None:
super().setUp()
self.log_records = []
self.listening_task: Optional[asyncio.Task] = None
self.mocking_assistant = NetworkMockingAssistant()
self.throttler = AsyncThrottler(CONSTANTS.RATE_LIMITS)
self.mock_time_provider = MagicMock()
self.mock_time_provider.time.return_value = 1000
self.auth = KucoinAuth(self.api_key,
self.api_passphrase,
self.api_secret_key,
time_provider=self.mock_time_provider)
self.time_synchronizer = TimeSynchronizer()
self.time_synchronizer.add_time_offset_ms_sample(0)
self.api_factory = web_utils.build_api_factory(
throttler=self.throttler,
time_synchronizer=self.time_synchronizer,
auth=self.auth)
self.data_source = KucoinAPIUserStreamDataSource(
throttler=self.throttler, api_factory=self.api_factory)
self.data_source.logger().setLevel(1)
self.data_source.logger().addHandler(self)
def setUp(self) -> None:
super().setUp()
self.throttler = AsyncThrottler(CONSTANTS.RATE_LIMITS)
self.auth = KucoinAuth(self.api_key, self.api_passphrase,
self.api_secret_key)
self.data_source = KucoinAPIUserStreamDataSource(
self.throttler, self.auth)
self.mocking_assistant = NetworkMockingAssistant()
def _create_user_stream_data_source(self) -> UserStreamTrackerDataSource:
return KucoinAPIUserStreamDataSource(
auth=self._auth,
trading_pairs=self._trading_pairs,
connector=self,
api_factory=self._web_assistants_factory,
domain=self.domain,
)
def __init__(self,
kucoin_api_key: str,
kucoin_passphrase: str,
kucoin_secret_key: str,
trading_pairs: Optional[List[str]] = None,
trading_required: bool = True,
domain: str = CONSTANTS.DEFAULT_DOMAIN):
self._domain = domain
self._time_synchronizer = TimeSynchronizer()
super().__init__()
self._auth = KucoinAuth(
api_key=kucoin_api_key,
passphrase=kucoin_passphrase,
secret_key=kucoin_secret_key,
time_provider=self._time_synchronizer)
self._throttler = AsyncThrottler(CONSTANTS.RATE_LIMITS)
self._api_factory = web_utils.build_api_factory(
throttler=self._throttler,
time_synchronizer=self._time_synchronizer,
domain=self._domain,
auth=self._auth)
self._last_poll_timestamp = 0
self._last_timestamp = 0
self._trading_pairs = trading_pairs
self._order_book_tracker = OrderBookTracker(
data_source=KucoinAPIOrderBookDataSource(
trading_pairs=trading_pairs,
domain=self._domain,
api_factory=self._api_factory,
throttler=self._throttler,
time_synchronizer=self._time_synchronizer),
trading_pairs=trading_pairs,
domain=self._domain)
self._user_stream_tracker = UserStreamTracker(
data_source=KucoinAPIUserStreamDataSource(
domain=self._domain,
api_factory=self._api_factory,
throttler=self._throttler))
self._poll_notifier = asyncio.Event()
self._status_polling_task = None
self._user_stream_tracker_task = None
self._user_stream_event_listener_task = None
self._trading_rules_polling_task = None
self._trading_fees_polling_task = None
self._trading_required = trading_required
self._trading_rules = {}
self._trading_fees = {}
self._order_tracker: ClientOrderTracker = ClientOrderTracker(connector=self)
class TestKucoinAPIUserStreamDataSource(unittest.TestCase):
@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.api_key = "someKey"
cls.api_passphrase = "somePassPhrase"
cls.api_secret_key = "someSecretKey"
def setUp(self) -> None:
super().setUp()
self.throttler = AsyncThrottler(CONSTANTS.RATE_LIMITS)
self.auth = KucoinAuth(self.api_key, self.api_passphrase,
self.api_secret_key)
self.data_source = KucoinAPIUserStreamDataSource(
self.throttler, self.auth)
self.mocking_assistant = NetworkMockingAssistant()
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
@staticmethod
def get_listen_key_mock():
listen_key = {
"code": "200000",
"data": {
"token":
"someToken",
"instanceServers": [{
"endpoint": "wss://someEndpoint",
"encrypt": True,
"protocol": "websocket",
"pingInterval": 18000,
"pingTimeout": 10000,
}]
}
}
return listen_key
@aioresponses()
def test_get_listen_key_raises(self, mock_api):
url = CONSTANTS.BASE_PATH_URL + CONSTANTS.PRIVATE_WS_DATA_PATH_URL
mock_api.post(url, status=500)
with self.assertRaises(IOError):
self.async_run_with_timeout(self.data_source.get_listen_key())
@aioresponses()
def test_get_listen_key(self, mock_api):
url = CONSTANTS.BASE_PATH_URL + CONSTANTS.PRIVATE_WS_DATA_PATH_URL
resp = self.get_listen_key_mock()
mock_api.post(url, body=json.dumps(resp))
ret = self.async_run_with_timeout(self.data_source.get_listen_key())
self.assertEqual(ret, resp) # shallow comparison ok
@aioresponses()
@patch("aiohttp.client.ClientSession.ws_connect", new_callable=AsyncMock)
def test_listen_to_user_stream_subscribes_to_private_topics(
self, mock_api, ws_connect_mock):
url = CONSTANTS.BASE_PATH_URL + CONSTANTS.PRIVATE_WS_DATA_PATH_URL
resp = self.get_listen_key_mock()
mock_api.post(url, body=json.dumps(resp))
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
msg_queue = asyncio.Queue()
self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop, msg_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(
ws_connect_mock.return_value)
self.assertEqual(len(CONSTANTS.PRIVATE_ENDPOINT_NAMES),
len(sent_messages))
subscribed_endpoints = {m["topic"] for m in sent_messages}
self.assertEqual(set(CONSTANTS.PRIVATE_ENDPOINT_NAMES),
subscribed_endpoints)
@aioresponses()
@patch("aiohttp.client.ClientSession.ws_connect", new_callable=AsyncMock)
def test_listen_to_user_stream_accepts_message(self, mock_api,
ws_connect_mock):
url = CONSTANTS.BASE_PATH_URL + CONSTANTS.PRIVATE_WS_DATA_PATH_URL
resp = self.get_listen_key_mock()
mock_api.post(url, body=json.dumps(resp))
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
msg = "someMsg"
msg_queue = asyncio.Queue()
self.mocking_assistant.add_websocket_aiohttp_message(
ws_connect_mock.return_value, json.dumps(msg))
self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop, msg_queue))
self.mocking_assistant.run_until_all_aiohttp_messages_delivered(
ws_connect_mock.return_value)
self.assertTrue(not msg_queue.empty())
queued = msg_queue.get_nowait()
self.assertEqual(msg, queued)
@aioresponses()
@patch("aiohttp.client.ClientSession.ws_connect", new_callable=AsyncMock)
def test_listen_to_user_stream_sends_ping_ignores_pong(
self, mock_api, ws_connect_mock):
url = CONSTANTS.BASE_PATH_URL + CONSTANTS.PRIVATE_WS_DATA_PATH_URL
resp = self.get_listen_key_mock()
mock_api.post(url, body=json.dumps(resp))
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
self.mocking_assistant.add_websocket_aiohttp_message(
ws_connect_mock.return_value,
message="",
message_type=aiohttp.WSMsgType.PONG)
msg = "someMsg"
self.mocking_assistant.add_websocket_aiohttp_message(
ws_connect_mock.return_value, json.dumps(msg))
msg_queue = asyncio.Queue()
self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop, msg_queue))
self.mocking_assistant.run_until_all_aiohttp_messages_delivered(
ws_connect_mock.return_value)
ws_connect_mock.return_value.ping.assert_called() # ping was sent
self.assertTrue(not msg_queue.empty())
queued = msg_queue.get_nowait()
self.assertEqual(msg, queued)
class TestKucoinAPIUserStreamDataSource(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.api_key = "someKey"
cls.api_passphrase = "somePassPhrase"
cls.api_secret_key = "someSecretKey"
def setUp(self) -> None:
super().setUp()
self.log_records = []
self.listening_task: Optional[asyncio.Task] = None
self.mocking_assistant = NetworkMockingAssistant()
self.throttler = AsyncThrottler(CONSTANTS.RATE_LIMITS)
self.mock_time_provider = MagicMock()
self.mock_time_provider.time.return_value = 1000
self.auth = KucoinAuth(
self.api_key,
self.api_passphrase,
self.api_secret_key,
time_provider=self.mock_time_provider)
client_config_map = ClientConfigAdapter(ClientConfigMap())
self.connector = KucoinExchange(
client_config_map=client_config_map,
kucoin_api_key="",
kucoin_passphrase="",
kucoin_secret_key="",
trading_pairs=[],
trading_required=False)
self.data_source = KucoinAPIUserStreamDataSource(
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)
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 async_run_with_timeout(self, coroutine: Awaitable, timeout: int = 1):
ret = self.ev_loop.run_until_complete(asyncio.wait_for(coroutine, timeout))
return ret
@staticmethod
def get_listen_key_mock():
listen_key = {
"code": "200000",
"data": {
"token": "someToken",
"instanceServers": [
{
"endpoint": "wss://someEndpoint",
"encrypt": True,
"protocol": "websocket",
"pingInterval": 18000,
"pingTimeout": 10000,
}
]
}
}
return listen_key
@aioresponses()
@patch("aiohttp.ClientSession.ws_connect", new_callable=AsyncMock)
@patch("hummingbot.connector.exchange.kucoin.kucoin_web_utils.next_message_id")
def test_listen_for_user_stream_subscribes_to_orders_and_balances_events(self, mock_api, id_mock, ws_connect_mock):
id_mock.side_effect = [1, 2]
url = web_utils.private_rest_url(path_url=CONSTANTS.PRIVATE_WS_DATA_PATH_URL)
resp = {
"code": "200000",
"data": {
"instanceServers": [
{
"endpoint": "wss://test.url/endpoint",
"protocol": "websocket",
"encrypt": True,
"pingInterval": 50000,
"pingTimeout": 10000
}
],
"token": "testToken"
}
}
mock_api.post(url, body=json.dumps(resp))
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock()
result_subscribe_trades = {
"type": "ack",
"id": 1
}
result_subscribe_diffs = {
"type": "ack",
"id": 2
}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=ws_connect_mock.return_value,
message=json.dumps(result_subscribe_trades))
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=ws_connect_mock.return_value,
message=json.dumps(result_subscribe_diffs))
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_subscription_messages = self.mocking_assistant.json_messages_sent_through_websocket(
websocket_mock=ws_connect_mock.return_value)
self.assertEqual(2, len(sent_subscription_messages))
expected_orders_subscription = {
"id": 1,
"type": "subscribe",
"topic": "/spotMarket/tradeOrders",
"privateChannel": True,
"response": False
}
self.assertEqual(expected_orders_subscription, sent_subscription_messages[0])
expected_balances_subscription = {
"id": 2,
"type": "subscribe",
"topic": "/account/balance",
"privateChannel": True,
"response": False
}
self.assertEqual(expected_balances_subscription, sent_subscription_messages[1])
self.assertTrue(self._is_logged(
"INFO",
"Subscribed to private order changes and balance updates channels..."
))
@aioresponses()
@patch("aiohttp.ClientSession.ws_connect", new_callable=AsyncMock)
def test_listen_for_user_stream_get_listen_key_successful_with_user_update_event(self, mock_api, mock_ws):
url = web_utils.private_rest_url(path_url=CONSTANTS.PRIVATE_WS_DATA_PATH_URL)
regex_url = re.compile(f"^{url}".replace(".", r"\.").replace("?", r"\?"))
mock_response = self.get_listen_key_mock()
mock_api.post(regex_url, body=json.dumps(mock_response))
mock_ws.return_value = self.mocking_assistant.create_websocket_mock()
order_event = {
"type": "message",
"topic": "/spotMarket/tradeOrders",
"subject": "orderChange",
"channelType": "private",
"data": {
"symbol": "KCS-USDT",
"orderType": "limit",
"side": "buy",
"orderId": "5efab07953bdea00089965d2",
"type": "open",
"orderTime": 1593487481683297666,
"size": "0.1",
"filledSize": "0",
"price": "0.937",
"clientOid": "1593487481000906",
"remainSize": "0.1",
"status": "open",
"ts": 1593487481683297666
}
}
self.mocking_assistant.add_websocket_aiohttp_message(mock_ws.return_value, 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)
)
msg = self.async_run_with_timeout(msg_queue.get())
self.assertEqual(order_event, msg)
mock_ws.return_value.ping.assert_called()
@aioresponses()
@patch("aiohttp.ClientSession.ws_connect", new_callable=AsyncMock)
def test_listen_for_user_stream_does_not_queue_pong_payload(self, mock_api, mock_ws):
url = web_utils.private_rest_url(path_url=CONSTANTS.PRIVATE_WS_DATA_PATH_URL)
regex_url = re.compile(f"^{url}".replace(".", r"\.").replace("?", r"\?"))
mock_response = self.get_listen_key_mock()
mock_pong = {
"id": "1545910590801",
"type": "pong"
}
mock_api.post(regex_url, body=json.dumps(mock_response))
mock_ws.return_value = self.mocking_assistant.create_websocket_mock()
self.mocking_assistant.add_websocket_aiohttp_message(mock_ws.return_value, json.dumps(mock_pong))
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())
@aioresponses()
@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, mock_api, sleep_mock, mock_ws):
url = web_utils.private_rest_url(path_url=CONSTANTS.PRIVATE_WS_DATA_PATH_URL)
regex_url = re.compile(f"^{url}".replace(".", r"\.").replace("?", r"\?"))
mock_response = self.get_listen_key_mock()
mock_api.post(regex_url, body=json.dumps(mock_response))
mock_ws.side_effect = Exception("TEST ERROR.")
sleep_mock.side_effect = asyncio.CancelledError # to finish the task execution
msg_queue = asyncio.Queue()
try:
self.async_run_with_timeout(self.data_source.listen_for_user_stream(msg_queue))
except asyncio.CancelledError:
pass
self.assertTrue(
self._is_logged("ERROR",
"Unexpected error while listening to user stream. Retrying after 5 seconds..."))
@aioresponses()
@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_iter_message_throws_exception(self, mock_api, sleep_mock, mock_ws):
url = web_utils.private_rest_url(path_url=CONSTANTS.PRIVATE_WS_DATA_PATH_URL)
regex_url = re.compile(f"^{url}".replace(".", r"\.").replace("?", r"\?"))
mock_response = self.get_listen_key_mock()
mock_api.post(regex_url, body=json.dumps(mock_response))
msg_queue: asyncio.Queue = asyncio.Queue()
mock_ws.return_value = self.mocking_assistant.create_websocket_mock()
mock_ws.return_value.receive.side_effect = Exception("TEST ERROR")
sleep_mock.side_effect = asyncio.CancelledError # to finish the task execution
try:
self.async_run_with_timeout(self.data_source.listen_for_user_stream(msg_queue))
except asyncio.CancelledError:
pass
self.assertTrue(
self._is_logged(
"ERROR",
"Unexpected error while listening to user stream. Retrying after 5 seconds..."))
@aioresponses()
@patch("aiohttp.ClientSession.ws_connect", new_callable=AsyncMock)
@patch("hummingbot.connector.exchange.kucoin.kucoin_web_utils.next_message_id")
@patch("hummingbot.connector.exchange.kucoin.kucoin_api_user_stream_data_source.KucoinAPIUserStreamDataSource"
"._time")
def test_listen_for_user_stream_sends_ping_message_before_ping_interval_finishes(
self,
mock_api,
time_mock,
id_mock,
ws_connect_mock):
id_mock.side_effect = [1, 2, 3, 4]
time_mock.side_effect = [1000, 1100, 1101, 1102] # Simulate first ping interval is already due
url = web_utils.private_rest_url(path_url=CONSTANTS.PRIVATE_WS_DATA_PATH_URL)
resp = {
"code": "200000",
"data": {
"instanceServers": [
{
"endpoint": "wss://test.url/endpoint",
"protocol": "websocket",
"encrypt": True,
"pingInterval": 20000,
"pingTimeout": 10000
}
],
"token": "testToken"
}
}
mock_api.post(url, body=json.dumps(resp))
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock()
result_subscribe_trades = {
"type": "ack",
"id": 1
}
result_subscribe_diffs = {
"type": "ack",
"id": 2
}
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=ws_connect_mock.return_value,
message=json.dumps(result_subscribe_trades))
self.mocking_assistant.add_websocket_aiohttp_message(
websocket_mock=ws_connect_mock.return_value,
message=json.dumps(result_subscribe_diffs))
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)
expected_ping_message = {
"id": 3,
"type": "ping",
}
self.assertEqual(expected_ping_message, sent_messages[-1])
def data_source(self) -> UserStreamTrackerDataSource:
if not self._data_source:
self._data_source = KucoinAPIUserStreamDataSource(
kucoin_auth=self._kucoin_client)
return self._data_source
class TestKucoinAPIUserStreamDataSource(unittest.TestCase):
@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.api_key = "someKey"
cls.api_passphrase = "somePassPhrase"
cls.api_secret_key = "someSecretKey"
def setUp(self) -> None:
super().setUp()
self.throttler = AsyncThrottler(CONSTANTS.RATE_LIMITS)
self.auth = KucoinAuth(self.api_key, self.api_passphrase,
self.api_secret_key)
self.data_source = KucoinAPIUserStreamDataSource(
self.throttler, self.auth)
self.mocking_assistant = NetworkMockingAssistant()
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
@staticmethod
def get_listen_key_mock():
listen_key = {
"code": "200000",
"data": {
"token":
"someToken",
"instanceServers": [{
"endpoint": "wss://someEndpoint",
"encrypt": True,
"protocol": "websocket",
"pingInterval": 18000,
"pingTimeout": 10000,
}]
}
}
return listen_key
@aioresponses()
def test_get_listen_key_raises(self, mock_api):
url = CONSTANTS.BASE_PATH_URL + CONSTANTS.PRIVATE_WS_DATA_PATH_URL
mock_api.post(url, status=500)
with self.assertRaises(IOError):
self.async_run_with_timeout(self.data_source.get_listen_key())
@aioresponses()
def test_get_listen_key(self, mock_api):
url = CONSTANTS.BASE_PATH_URL + CONSTANTS.PRIVATE_WS_DATA_PATH_URL
resp = self.get_listen_key_mock()
mock_api.post(url, body=json.dumps(resp))
ret = self.async_run_with_timeout(self.data_source.get_listen_key())
self.assertEqual(ret, resp) # shallow comparison ok
@aioresponses()
@patch("websockets.connect", new_callable=AsyncMock)
def test_listen_to_user_stream(self, mock_api, ws_connect_mock):
url = CONSTANTS.BASE_PATH_URL + CONSTANTS.PRIVATE_WS_DATA_PATH_URL
resp = self.get_listen_key_mock()
mock_api.post(url, body=json.dumps(resp))
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
msg = "someMsg"
msg_queue = asyncio.Queue()
self.mocking_assistant.add_websocket_text_message(
ws_connect_mock.return_value, json.dumps(msg))
self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop, msg_queue))
self.mocking_assistant.run_until_all_text_messages_delivered(
ws_connect_mock.return_value)
self.assertTrue(not msg_queue.empty())
queued = msg_queue.get_nowait()
self.assertEqual(msg, queued)
@aioresponses()
@patch("websockets.connect", new_callable=AsyncMock)
def test_listen_for_user_stream_closes_ws_on_exception(
self, mock_api, ws_connect_mock):
url = CONSTANTS.BASE_PATH_URL + CONSTANTS.PRIVATE_WS_DATA_PATH_URL
resp = self.get_listen_key_mock()
mock_api.post(url, body=json.dumps(resp))
raised_event = asyncio.Event()
async def raise_exception():
raised_event.set()
raise IOError
ws_connect_mock.return_value = self.mocking_assistant.create_websocket_mock(
)
ws_connect_mock.return_value.recv.side_effect = raise_exception
self.ev_loop.create_task(
self.data_source.listen_for_user_stream(self.ev_loop,
asyncio.Queue()))
self.async_run_with_timeout(coroutine=raised_event.wait())
ws_connect_mock.return_value.close.assert_called()