async def ws_auth(self) -> Dict[Any, Any]: ws = GateIoWebsocket(self.auth) await ws.connect() await ws.subscribe(Constants.WS_SUB["USER_BALANCE"]) async for response in ws.on_message(): if ws.is_subscribed: return True return False
async def ws_auth(self) -> Dict[Any, Any]: ws = GateIoWebsocket(self.auth) await ws.connect() await ws.subscribe(CONSTANTS.USER_BALANCE_ENDPOINT_NAME) async for response in ws.on_message(): if ws.is_subscribed: return True return False
async def ws_auth(self) -> Dict[Any, Any]: ws = GateIoWebsocket(api_factory=build_gate_io_api_factory( throttler=AsyncThrottler(CONSTANTS.RATE_LIMITS)), auth=self.auth) await ws.connect() await ws.subscribe(CONSTANTS.USER_BALANCE_ENDPOINT_NAME) async for response in ws.on_message(): if ws.is_subscribed: return True return False
def setUp(self) -> None: self.ev_loop = asyncio.get_event_loop() self.ws = GateIoWebsocket() self.mocking_assistant = NetworkMockingAssistant()
class GateIoWebsocketTest(unittest.TestCase): def setUp(self) -> None: self.ev_loop = asyncio.get_event_loop() self.ws = GateIoWebsocket() self.mocking_assistant = NetworkMockingAssistant() 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) def test_subscribe(self, mock_ws): subscription_channel = "someChannel" mock_ws.return_value = self.mocking_assistant.create_websocket_mock() self.async_run_with_timeout(self.ws.connect()) ret = self.async_run_with_timeout(self.ws.subscribe(channel=subscription_channel)) np.testing.assert_allclose([time.time()], [ret], rtol=1) calls = self.mocking_assistant.json_messages_sent_through_websocket(mock_ws.return_value) self.assertEqual(1, len(calls)) subscription_call = calls[0] self.assertEqual("subscribe", subscription_call["event"]) self.assertEqual(ret, subscription_call["time"]) self.assertEqual(subscription_channel, subscription_call["channel"]) @patch("aiohttp.client.ClientSession.ws_connect", new_callable=AsyncMock) def test_unsubscribe(self, mock_ws): subscription_channel = "someChannel" mock_ws.return_value = self.mocking_assistant.create_websocket_mock() self.async_run_with_timeout(self.ws.connect()) ret = self.async_run_with_timeout(self.ws.unsubscribe(channel=subscription_channel)) np.testing.assert_allclose([time.time()], [ret], rtol=1) calls = self.mocking_assistant.json_messages_sent_through_websocket(mock_ws.return_value) self.assertEqual(1, len(calls)) subscription_call = calls[0] self.assertEqual("unsubscribe", subscription_call["event"]) self.assertEqual(ret, subscription_call["time"]) self.assertEqual(subscription_channel, subscription_call["channel"]) @patch("aiohttp.client.ClientSession.ws_connect", new_callable=AsyncMock) def test_on_message(self, mock_ws): mock_ws.return_value = self.mocking_assistant.create_websocket_mock() self.async_run_with_timeout(self.ws.connect()) async_iter = self.ws.on_message() mock_event = "somEvent" self.mocking_assistant.add_websocket_aiohttp_message( mock_ws.return_value, json.dumps({"event": mock_event}) ) ret = self.async_run_with_timeout(async_iter.__anext__(), timeout=0.1) self.assertEqual(mock_event, ret["event"]) @patch("aiohttp.client.ClientSession.ws_connect", new_callable=AsyncMock) def test_ping_sent_pong_ignored(self, mock_ws): mock_ws.return_value = self.mocking_assistant.create_websocket_mock() mock_ws.return_value.closed = False self.async_run_with_timeout(self.ws.connect()) async_iter = self.ws.on_message() self.mocking_assistant.add_websocket_aiohttp_message( mock_ws.return_value, message="", message_type=aiohttp.WSMsgType.PONG # should be ignored ) mock_event = "somEvent" self.mocking_assistant.add_websocket_aiohttp_message( mock_ws.return_value, json.dumps({"event": mock_event}) ) ret = self.async_run_with_timeout(async_iter.__anext__(), timeout=1) mock_ws.return_value.ping.assert_called() self.assertEqual(mock_event, ret["event"]) @patch("aiohttp.client.ClientSession.ws_connect", new_callable=AsyncMock) def test_last_recv_time_set_on_pong(self, mock_ws): mock_ws.return_value = self.mocking_assistant.create_websocket_mock() mock_ws.return_value.closed = False self.async_run_with_timeout(self.ws.connect()) async_iter = self.ws.on_message() self.mocking_assistant.add_websocket_aiohttp_message( mock_ws.return_value, message="", message_type=aiohttp.WSMsgType.PONG # should be ignored ) anext_task = self.ev_loop.create_task(async_iter.__anext__()) try: self.mocking_assistant.run_until_all_aiohttp_messages_delivered(mock_ws.return_value) np.testing.assert_allclose([self.ws.last_recv_time], [time.time()], rtol=1) except Exception: raise finally: anext_task.cancel() @patch("aiohttp.client.ClientSession.ws_connect", new_callable=AsyncMock) def test_backup_ping_pong(self, mock_ws): mock_ws.return_value = self.mocking_assistant.create_websocket_mock() mock_ws.return_value.closed = False default_receive = mock_ws.return_value.receive.side_effect self.async_run_with_timeout(self.ws.connect()) async_iter = self.ws.on_message() async def switch_back_and_raise_timeout(*args, **kwargs): mock_ws.return_value.receive.side_effect = default_receive raise asyncio.TimeoutError mock_ws.return_value.receive.side_effect = switch_back_and_raise_timeout self.mocking_assistant.add_websocket_aiohttp_message( mock_ws.return_value, json.dumps({"channel": "spot.pong"}) # should be ignored ) mock_event = "somEvent" self.mocking_assistant.add_websocket_aiohttp_message( mock_ws.return_value, json.dumps({"event": mock_event}) ) ret = self.async_run_with_timeout(async_iter.__anext__(), timeout=1) self.assertEqual(mock_event, ret["event"])