diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 7463b4cb..24773b56 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -285,7 +285,7 @@ def subscribe(self, *subscriptions: str): # If user subs to X.*, remove other X.\w+ if sym == "*": for t in list(self.subs): - if t.startswith(topic): + if t != s and t.strip().split(".", 1)[0] == topic: self.scheduled_subs.discard(t) self.schedule_resub = True @@ -306,7 +306,7 @@ def unsubscribe(self, *subscriptions: str): # If user unsubs to X.*, remove other X.\w+ if sym == "*": for t in list(self.subs): - if t.startswith(topic): + if t.strip().split(".", 1)[0] == topic: self.scheduled_subs.discard(t) self.schedule_resub = True diff --git a/test_websocket/test_wildcard.py b/test_websocket/test_wildcard.py new file mode 100644 index 00000000..9224a165 --- /dev/null +++ b/test_websocket/test_wildcard.py @@ -0,0 +1,28 @@ +import unittest +from massive import WebSocketClient + + +class WildcardTest(unittest.TestCase): + def client(self, *subs): + c = WebSocketClient(api_key="", subscriptions=list(subs)) + c.subs = set(subs) # as if connect() had sent them + return c + + def test_subscribe_wildcard_keeps_other_channels(self): + c = self.client("A.MSFT", "AM.AAPL") + c.subscribe("A.*") + self.assertEqual(c.scheduled_subs, {"A.*", "AM.AAPL"}) + + def test_subscribe_wildcard_twice_keeps_it(self): + c = self.client("A.*") + c.subscribe("A.*") + self.assertEqual(c.scheduled_subs, {"A.*"}) + + def test_unsubscribe_wildcard_keeps_other_channels(self): + c = self.client("A.*", "A.MSFT", "AM.AAPL") + c.unsubscribe("A.*") + self.assertEqual(c.scheduled_subs, {"AM.AAPL"}) + + +if __name__ == "__main__": + unittest.main()