enable multiple calls to subscribe for same event type (#114)

* enable multiple calls to subscribe for same event type * fix streamer test * fix lint
tastyware · Dec 6, 2023 · 5be6413 · 5be6413
1 parent 09fad18
commit 5be6413
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 11 deletions.
diff --git a/tastytrade/streamer.py b/tastytrade/streamer.py
@@ -678,7 +678,6 @@ async def _connect(self) -> None:
         Connect to the websocket server using the URL and
         authorization token provided during initialization.
         """
-
         async with websockets.connect(  # type: ignore
             self._wss_url
         ) as websocket:
@@ -701,8 +700,7 @@ async def _connect(self) -> None:
                 elif message['type'] == 'CHANNEL_OPENED':
                     channel = next((k for k, v in self._channels.items()
                                     if v == message['channel']))
-                    self._subscription_state[channel] \
-                        = message['type']
+                    self._subscription_state[channel] = message['type']
                 elif message['type'] == 'CHANNEL_CLOSED':
                     pass
                 elif message['type'] == 'FEED_CONFIG':
@@ -712,7 +710,7 @@ async def _connect(self) -> None:
                 elif message['type'] == 'KEEPALIVE':
                     pass
                 else:
-                    raise TastytradeError(message)
+                    raise TastytradeError('Unknown message type:', message)
 
     async def _setup_connection(self):
         message = {
@@ -782,12 +780,12 @@ async def subscribe(
         :param event_type: type of subscription to add
         :param symbols: list of symbols to subscribe for
         """
-        await self._channel_request(event_type)
-        event_type_str = str(event_type).split('.')[1].capitalize()
+        if self._subscription_state[event_type] != 'CHANNEL_OPENED':
+            await self._channel_request(event_type)
         message = {
             'type': 'FEED_SUBSCRIPTION',
             'channel': self._channels[event_type],
-            'add': [{'symbol': symbol, "type": event_type_str}
+            'add': [{'symbol': symbol, 'type': event_type}
                     for symbol in symbols]
         }
         logger.debug('sending subscription: %s', message)

diff --git a/tests/test_streamer.py b/tests/test_streamer.py
@@ -20,8 +20,6 @@ async def test_account_streamer(session):
 
 @pytest.mark.asyncio
 async def test_dxlink_streamer(session):
-    message = "[{'eventType': 'Quote', 'eventSymbol': 'SPY', 'eventTime': 0, 'sequence': 0, 'timeNanoPart': 0, 'bidTime': 0, 'bidExchangeCode': 'Q', 'bidPrice': 450.5, 'bidSize': 796.0, 'askTime': 0, 'askExchangeCode': 'Q', 'askPrice': 450.55, 'askSize': 1100.0}, {'eventType': 'Quote', 'eventSymbol': 'AAPL', 'eventTime': 0, 'sequence': 0, 'timeNanoPart': 0, 'bidTime': 0, 'bidExchangeCode': 'Q', 'bidPrice': 190.39, 'bidSize': 1.0, 'askTime': 0, 'askExchangeCode': 'Q', 'askPrice': 190.44, 'askSize': 3.0}]"  # noqa: E501
-
     async with DXLinkStreamer(session) as streamer:
         subs = ['SPY', 'AAPL']
         await streamer.subscribe(EventType.QUOTE, subs)
@@ -32,5 +30,3 @@ async def test_dxlink_streamer(session):
             break
         await streamer.unsubscribe_candle(subs[0], '1d')
         await streamer.unsubscribe(EventType.QUOTE, subs[1])
-
-        streamer._map_message(message)