def test_process_clear_action_response(sm: SubscriptionManager):
category = Category.MARKET
events = [MarketEvents.ORDERS, MarketEvents.CANCELS]
topics = ['ETH_AURA', 'ETH_ZRX']
payload = dict(events=events, topics=topics)
sm._process_subscribe_action_response(category, payload)
assert len(sm.subscriptions[category].topics) == len(topics)
sm._process_clear_action_response(category, payload)
assert len(sm.subscriptions[category].topics) == 0
async def test_resubscribe(sm: SubscriptionManager):
sub = Subscription(category=Category.MARKET,
events=[MarketEvents.CANCELS, MarketEvents.ORDERS],
topics=['ETH_AURA', 'ETH_ZRX'])
sm.subscriptions = {Category.MARKET: sub}
sm._init_subscriptions = Mock()
sm.subscribe = CoroutineMock()
await sm.resubscribe()
sm._init_subscriptions.assert_called_once()
sm.subscribe.assert_awaited_once_with(sub)
def test_sub_payload(sm: SubscriptionManager):
sm._filter_none = Mock(return_value='value')
params = dict(action=Action.SUBSCRIBE,
topics=['ETH_AURA'],
events=MarketEvents.ORDERS)
result = sm._sub_payload(**params)
sm._filter_none.assert_called_once_with(
dict(action=Action.SUBSCRIBE.value,
topics=['ETH_AURA'],
events=MarketEvents.ORDERS))
assert result == 'value'
def test_process_sub_response_get_handler(sm: SubscriptionManager):
message = dict(
result='success',
request='subscribeToMarkets',
payload=dict(action='get'),
)
sm._logger.error = Mock()
sm._process_get_action_response = Mock()
result = sm.process_sub_response(message)
sm._process_get_action_response.assert_called_once_with(
Category(message['request']), message['payload'])
assert result is None
async def test_subscribe_warning(sm: SubscriptionManager):
sm._logger.warning = Mock()
send_mock = CoroutineMock()
send_mock.return_value = 'some return value'
sm._ds.send_message = send_mock
sm._sub_payload = Mock()
sub_payload_value = 'some value'
sm._sub_payload.return_value = sub_payload_value
sub = Subscription(category=Category.MARKET,
events=[MarketEvents.CANCELS, MarketEvents.ORDERS],
topics=['ETH_AURA', 'ETH_ZRX'])
sm.subscriptions[sub.category] = {}
result = await sm.subscribe(sub)
sm._logger.warning.assert_called_once()
sm._sub_payload.assert_called_once_with(Action.SUBSCRIBE,
topics=sub.topics,
events=sub.events)
send_mock.assert_awaited_once_with(sub.category.value, sub_payload_value,
None)
assert result == 'some return value'
def test_process_subscribe_action_response(sm: SubscriptionManager):
category = Category.MARKET
events = [MarketEvents.ORDERS, MarketEvents.CANCELS]
topics = ['ETH_AURA', 'ETH_ZRX']
payload = dict(events=events, topics=topics)
assert sm.subscriptions == {}
sm._process_subscribe_action_response(category, payload)
assert category in sm.subscriptions
sub = sm.subscriptions[category]
assert sub.category == category
assert len(sub.events) == len(events)
assert sub.events == (MarketEvents.CANCELS.value,
MarketEvents.ORDERS.value)
assert len(sub.topics) == len(topics)
assert sub.topics == ('ETH_AURA', 'ETH_ZRX')
def __init__(
self,
api_key: str = '17paIsICur8sA0OBqG6dH5G1rmrHNMwt4oNk4iX9',
version: str = '1.0.0',
ws_endpoint: str = 'wss://datastream.idex.market',
handshake_timeout: float = 1.0,
return_sub_responses=False,
loop: AbstractEventLoop = None
):
self._API_KEY = api_key
self._WS_ENDPOINT = ws_endpoint
self._WS_VERSION = version
self._HANDSHAKE_TIMEOUT = handshake_timeout
self._loop = loop or asyncio.get_event_loop()
self._logger = logging.getLogger(__name__)
self._rid = ShortId()
self.sub_manager = SubscriptionManager(self, return_sub_responses)
async def test_get(sm: SubscriptionManager):
sm._ds.send_message = CoroutineMock(return_value='some return value')
sm._sub_payload = Mock(return_value='some value')
category = Category.MARKET
result = await sm.get(category)
sm._sub_payload.assert_called_once_with(Action.GET)
sm._ds.send_message.assert_awaited_once_with(category.value, 'some value',
None)
assert result == 'some return value'
async def test_init():
datastream = Mock()
return_sub_responses = False
with patch(
'aioidex.datastream.sub_manager.SubscriptionManager._init_subscriptions'
) as mock:
sm = SubscriptionManager(datastream, return_sub_responses)
mock.assert_called_once()
assert sm._ds is datastream
assert sm._return_responses == return_sub_responses
assert isinstance(sm._logger, Logger)
assert sm._CATEGORY_VALUES == set(_.value for _ in Category)
async def test_clear_rid(sm: SubscriptionManager):
sm._ds.send_message = CoroutineMock(return_value='some return value')
sm._sub_payload = Mock(return_value='some value')
category = Category.MARKET
result = await sm.clear(category, 'some_rid')
sm._sub_payload.assert_called_once_with(Action.CLEAR)
sm._ds.send_message.assert_awaited_once_with(category.value, 'some value',
'some_rid')
assert result == 'some return value'
async def test_unsubscribe_rid(sm: SubscriptionManager):
sm._ds.send_message = CoroutineMock(return_value='some return value')
sm._sub_payload = Mock(return_value='some value')
category = Category.MARKET
result = await sm.unsubscribe(category, ['ETH_AURA'], 'some_rid')
sm._sub_payload.assert_called_once_with(Action.UNSUBSCRIBE, ['ETH_AURA'])
sm._ds.send_message.assert_awaited_once_with(category.value, 'some value',
'some_rid')
assert result == 'some return value'
def test_process_sub_response_error(sm: SubscriptionManager):
message = dict(
result='error',
request='subscribeToMarkets',
payload=dict(action='get'),
)
sm._logger.error = Mock()
result = sm.process_sub_response(message)
sm._logger.error.assert_called_once_with('Subscription error: %s', message)
assert result is None
def test_process_sub_response_no_handler(sm: SubscriptionManager):
message = dict(
result='success',
request='subscribeToMarkets',
payload=dict(action='unknown'),
)
sm._logger.error = Mock()
result = sm.process_sub_response(message)
sm._logger.error.assert_called_once_with(
'Handler for the subscription response not found: %s', message)
assert result is None
def test_process_unsubscribe_action_response(sm: SubscriptionManager):
category = Category.MARKET
events = [MarketEvents.ORDERS, MarketEvents.CANCELS]
topics = ['ETH_AURA', 'ETH_ZRX']
payload = dict(events=events, topics=topics)
sm._process_subscribe_action_response(category, payload)
old_topics = sm.subscriptions[category].topics
sm._process_unsubscribe_action_response(Category.CHAIN, payload)
assert old_topics == sm.subscriptions[category].topics
topics = ['ETH_AURA']
payload = dict(events=events, topics=topics)
sm._process_unsubscribe_action_response(category, payload)
assert sm.subscriptions[category].topics == ('ETH_AURA', )
def test_process_get_action_response(sm: SubscriptionManager):
category = Category.MARKET
events = [MarketEvents.ORDERS, MarketEvents.CANCELS]
topics = ['ETH_AURA', 'ETH_ZRX']
payload = dict(events=events, topics=topics)
assert sm.subscriptions == {}
sm._process_get_action_response(category, payload)
assert sm.subscriptions == {}
sm._process_subscribe_action_response(category, payload)
payload['topics'] = ['ETH_SAN']
old_events = sm.subscriptions[category].events
sm._process_get_action_response(category, payload)
assert sm.subscriptions[category].events == old_events
assert sm.subscriptions[category].topics == ('ETH_SAN', )
class IdexDatastream:
_ws: WebSocketClientProtocol = None
_sid: str = None
def __init__(
self,
api_key: str = '17paIsICur8sA0OBqG6dH5G1rmrHNMwt4oNk4iX9',
version: str = '1.0.0',
ws_endpoint: str = 'wss://datastream.idex.market',
handshake_timeout: float = 1.0,
return_sub_responses=False,
loop: AbstractEventLoop = None
):
self._API_KEY = api_key
self._WS_ENDPOINT = ws_endpoint
self._WS_VERSION = version
self._HANDSHAKE_TIMEOUT = handshake_timeout
self._loop = loop or asyncio.get_event_loop()
self._logger = logging.getLogger(__name__)
self._rid = ShortId()
self.sub_manager = SubscriptionManager(self, return_sub_responses)
async def _ping_ws_task(self, delay=30):
while True:
await asyncio.sleep(delay)
try:
await self._ws.ping()
except Exception as e:
self._logger.error('Ping task exception (%s): %s', type(e).__name__, e)
self._logger.warning('Reconnecting...')
await self.init()
await self.sub_manager.resubscribe()
async def _check_connection(self):
if not self._ws:
self._logger.info('Connection not created yet, creating...')
await self.init()
async def init(self, ws: WebSocketClientProtocol = None):
await self._init_connection(ws)
await self._shake_hand()
@backoff.on_exception(backoff.expo, Exception, max_time=30) # TODO: clarify exception
async def _init_connection(self, ws: WebSocketClientProtocol = None):
self._ws = ws or await self.create_connection()
self._logger.info('WS connection created: %s, %s', self._ws, self._ws.state)
async def create_connection(self):
return await websockets.connect(self._WS_ENDPOINT)
async def send_message(self, request: str, payload: Dict, rid: str = None) -> str:
await self._check_connection()
request_rid = rid or self._get_rid()
message = self._compose_message(request_rid, request, payload)
await self._ws.send(self._encode(message))
self._logger.debug('Sent message: %s', message)
return request_rid
def _compose_message(self, rid: str, request: str, payload: Dict):
return dict(
rid=rid,
sid=self._sid,
request=request,
payload=payload
)
async def listen(self):
await self._check_connection()
asyncio.create_task(self._ping_ws_task())
while True:
try:
# 1000 and 1001 exit codes reconnect support
if self._ws.closed:
raise websockets.ConnectionClosed(self._ws.close_code, 'Connection is closed')
async for msg in self._ws:
self._logger.debug('New message: %s', msg)
message = self._process_message(msg)
if message:
yield message
except (websockets.ConnectionClosed, IdexResponseSidError) as e:
self._logger.error(e)
self._logger.warning('Reconnecting...')
await self.init()
await self.sub_manager.resubscribe()
def _process_message(self, message: str) -> Optional[Dict]:
decoded_msg = self._decode(message)
self._logger.debug('New message: %s', decoded_msg)
self._check_warnings(decoded_msg)
self._check_errors(decoded_msg)
self._check_sid(decoded_msg)
if self.sub_manager.is_sub_response(decoded_msg):
return self.sub_manager.process_sub_response(decoded_msg)
return decoded_msg
def _check_warnings(self, message: Dict):
'''When an upcoming change to the specification is expected, a handshake request may return a warnings property
which can be used to alert you when an update to your integration will be required.
https://docs.idex.market/#tag/Datastream-Versioning
'''
if 'warnings' in message:
for warning in message['warnings']:
self._logger.warning(f'Response warning: {warning}')
@staticmethod
def _check_errors(message: Dict):
if message.get('result') == 'error':
payload = message['payload']
error_msg = payload.get('message', message)
raise IdexDataStreamError(f'Response error: {error_msg}')
def _check_sid(self, message: Dict):
'''Your client should monitor the sid value with every response and immediately reconnect if not a match.'''
sid = message['sid']
if sid != self._sid:
raise IdexResponseSidError(f'Received sid {sid!r} differs from existing sid {self._sid!r}, reconnecting...')
async def _shake_hand(self):
self._set_sid(None)
self._logger.info('Shaking hand...')
await self.send_message('handshake', dict(version=self._WS_VERSION, key=self._API_KEY))
self._process_handshake_response(await self._wait_for_handshake_response())
async def _wait_for_handshake_response(self) -> Dict:
try:
logging.info('Waiting for handshake response with timeout %s seconds...', self._HANDSHAKE_TIMEOUT)
response = await asyncio.wait_for(self._ws.recv(), self._HANDSHAKE_TIMEOUT, loop=self._loop)
except websockets.ConnectionClosed as e:
if e.code == 1002:
if e.reason == 'AuthenticationFailure':
raise IdexAuthenticationFailure(e)
elif e.reason == 'InvalidVersion':
raise IdexInvalidVersion(e)
raise
except asyncio.TimeoutError:
raise IdexHandshakeTimeout(f'Handshake response is not received within {self._HANDSHAKE_TIMEOUT} seconds')
else:
return self._decode(response)
def _process_handshake_response(self, message: Dict):
if message['result'] != 'success' or message['request'] != 'handshake':
raise IdexHandshakeException(message)
self._logger.info('Got handshake response: %s', message)
self._set_sid(message['sid'])
def _get_rid(self) -> str:
return f'rid:{self._rid.generate()}'
def _set_sid(self, sid: Union[str, None]):
if self._sid == sid:
return
self._logger.info('Sid changed from %r to %r', self._sid, sid)
self._sid = sid
@staticmethod
def _encode(data: Dict) -> str:
return ujson.dumps(data)
@staticmethod
def _decode(data: str) -> Dict:
decoded_msg = ujson.loads(data)
for field in ('payload', 'warnings'):
if field in decoded_msg and isinstance(decoded_msg[field], str):
decoded_msg[field] = ujson.loads(decoded_msg[field])
return decoded_msg
def test_filter_none(sm: SubscriptionManager):
assert sm._filter_none({1: 2}) == {1: 2}
assert sm._filter_none({1: 2, 3: None}) == {1: 2}
def test_init_subscriptions(sm: SubscriptionManager):
sm._init_subscriptions()
assert isinstance(sm.subscriptions, dict)
assert sm.subscriptions == {}
def test_is_sub_response(sm: SubscriptionManager):
assert sm.is_sub_response({'request': Category.MARKET.value}) is True
assert sm.is_sub_response({'request': 'not sub'}) is False