Fix streaming test

Author: ttt733

alpaca_trade_api/polygon/stream2.py

@@ -174,16 +174,16 @@ async def _dispatch(self, channel, msg):
                     ent = self._cast(channel, msg)
                     await handler(self, channel, ent)
 
-    def register(self, channel_pat, func):
+    def register(self, channel_pat, func, symbols=None):
         if not asyncio.iscoroutinefunction(func):
             raise ValueError('handler must be a coroutine function')
         if isinstance(channel_pat, str):
             channel_pat = re.compile(channel_pat)
         self._handlers[channel_pat] = func
+        self._handler_symbols[func] = symbols
 
     def deregister(self, channel_pat):
         if isinstance(channel_pat, str):
             channel_pat = re.compile(channel_pat)
-        handler = self._handlers[channel_pat]
-        del self._handler_symbols[handler]
+        self._handler_symbols.pop(self._handlers[channel_pat], None)
         del self._handlers[channel_pat]

alpaca_trade_api/stream2.py

@@ -131,27 +131,25 @@ async def _dispatch(self, channel, msg):
 
     def on(self, channel_pat, symbols=None):
         def decorator(func):
-            self.register(channel_pat, func)
-            if symbols:
-                self._handler_symbols[func] = symbols
+            self.register(channel_pat, func, symbols)
             return func
 
         return decorator
 
-    def register(self, channel_pat, func):
+    def register(self, channel_pat, func, symbols=None):
         if not asyncio.iscoroutinefunction(func):
             raise ValueError('handler must be a coroutine function')
         if isinstance(channel_pat, str):
             channel_pat = re.compile(channel_pat)
         self._handlers[channel_pat] = func
+        self._handler_symbols[func] = symbols
         if self.polygon:
-            self.polygon.register(channel_pat, func)
+            self.polygon.register(channel_pat, func, symbols)
 
     def deregister(self, channel_pat):
         if isinstance(channel_pat, str):
             channel_pat = re.compile(channel_pat)
-        handler = self._handlers[channel_pat]
-        del self._handler_symbols[handler]
+        self._handler_symbols.pop(self._handlers[channel_pat], None)
         del self._handlers[channel_pat]
         if self.polygon:
             self.polygon.deregister(channel_pat)

tests/test_stream2.py

@@ -1,4 +1,5 @@
 from alpaca_trade_api.stream2 import StreamConn
+from alpaca_trade_api.polygon.stream2 import StreamConn as PolyStream
 from alpaca_trade_api.entity import Account
 from alpaca_trade_api.polygon.entity import Entity as PolyEntity
 import asyncio
@@ -126,12 +127,15 @@ async def on_raise(conn, stream, msg):
     ent = conn._cast('other', {'key': 'value'})
     assert ent.key == 'value'
 
-    # _dispatch_nats
+    # polygon _dispatch
     conn = StreamConn('key-id', 'secret-key')
+    conn.polygon = PolyStream('key-id')
+    msg_data = PolyEntity({'key': 'value'})
+    conn.polygon._cast = mock.Mock(return_value=msg_data)
 
-    @conn.on('^Q.')
+    @conn.on('Q')
     async def on_q(conn, subject, data):
         on_q.data = data
 
-    _run(conn._dispatch_polygon(conn, 'Q.SPY', PolyEntity({'key': 'value'})))
+    _run(conn.polygon._dispatch('Q', msg_data))
     assert on_q.data.key == 'value'