fix: keep pending WebSocket subscriptions across a reconnect

massive/websocket/__init__.py

@@ -103,6 +103,9 @@ async def connect(
         :param close_timeout: How long to wait for handshake when calling .close.
         :raises AuthError: If invalid API key is supplied.
         """
+        # a new connection has no subscriptions, whatever an earlier call left
+        self.subs = set()
+        self.schedule_resub = True
         reconnects = 0
         logger.debug("connect: %s", self.url)
         ssl_context = self._ssl_context()
@@ -175,7 +178,8 @@ async def connect(
                 logger.debug("connection closed (ERR): %s", e)
                 self._flush_server_errors(force=True)
                 reconnects += 1
-                self.scheduled_subs = set(self.subs)
+                # keep scheduled_subs, which may hold changes not sent yet,
+                # and subscribe to all of it on the next connection
                 self.subs = set()
                 self.schedule_resub = True
                 if self.max_reconnects is not None and reconnects > self.max_reconnects:

test_websocket/base_ws.py

@@ -1,8 +1,9 @@
 import unittest
 import asyncio
-from mock_server import run_mock_server
+from mock_server import run_mock_server, inbox, port
 from typing import List
 from massive.websocket import WebSocketMessage
+from websockets.asyncio.client import connect
 
 unittest.util._MAX_LENGTH = 30000  # type: ignore
 
@@ -23,6 +24,31 @@ async def asyncSetUp(self):
         self.maxDiff = None
         loop = asyncio.get_event_loop()
         self.task = loop.create_task(run_mock_server())
+        inbox.clear()  # entries left by earlier tests
+        # wait until the mock server is accepting connections
+        for _ in range(100):
+            try:
+                async with connect(f"ws://localhost:{port}"):
+                    break
+            except OSError:
+                await asyncio.sleep(0.05)
+        else:
+            self.fail("mock server did not start")
+        await self.until(lambda: inbox)  # the connection above
+        inbox.clear()
 
     async def asyncTearDown(self):
+        # wait for the server to shut down so the next test can bind the port
         self.task.cancel()
+        try:
+            await self.task
+        except asyncio.CancelledError:
+            pass
+
+    async def until(self, condition):
+        # wait for condition() to hold, checking every 10 ms for up to 5 s
+        for _ in range(500):
+            if condition():
+                return
+            await asyncio.sleep(0.01)
+        self.fail("timed out")

test_websocket/mock_server.py

@@ -1,18 +1,23 @@
 from websockets import serve  # type: ignore
 import asyncio
 import json
-from typing import Set
+from typing import Any, List, Set
 
 port = 8765
 subs: Set[str] = set()
+# the messages each connection received, in order; one list per connection
+inbox: List[List[Any]] = []
 
 
 async def mock_server(websocket):
+    received: List[Any] = []
+    inbox.append(received)
     await websocket.send(
         '[{"ev":"status","status":"connected","message":"Connected Successfully"}]'
     )
     async for message in websocket:
         message = json.loads(message)
+        received.append(message)
         if "action" not in message:
             continue
         if message["action"] == "auth" and message["params"] == "MAX-CONNECTIONS":
@@ -35,6 +40,13 @@ async def mock_server(websocket):
                 await websocket.close(code=1000)
                 return
             for p in message["params"].split(","):
+                if p == "DROP.X":
+                    # emulate the connection failing, right after a trade
+                    await websocket.send(
+                        '[{"ev":"T","sym":"X","i":"1","x":1,"p":1,"s":1,"t":1,"q":1}]'
+                    )
+                    await websocket.close(code=1011)
+                    return
                 if p.startswith("ERR."):
                     # emulate the server rejecting a subscription
                     await websocket.send(

test_websocket/test_reconnect.py

@@ -0,0 +1,45 @@
+from massive import WebSocketClient
+from base_ws import BaseTest
+from mock_server import inbox, port
+import asyncio
+
+
+def sent(conn, action):
+    # the topics in each `action` message one connection received
+    return [set(m["params"].split(",")) for m in conn if m.get("action") == action]
+
+
+class ReconnectTest(BaseTest):
+    async def test_changes_made_before_a_drop_are_kept(self):
+        c = WebSocketClient(
+            api_key="",
+            feed=f"localhost:{port}",
+            secure=False,
+            subscriptions=["T.AAPL", "T.MSFT", "DROP.X"],
+        )
+
+        async def processor(msgs):
+            if len(inbox) == 1:
+                # DROP.X makes the server drop the connection; change the
+                # subscriptions before the client sends anything else
+                await c.websocket.wait_closed()
+                c.unsubscribe("T.MSFT", "DROP.X")
+                c.subscribe("T.NVDA")
+
+        task = asyncio.create_task(c.connect(processor))
+        await self.until(lambda: len(inbox) == 2 and sent(inbox[1], "subscribe"))
+        self.assertEqual(sent(inbox[1], "subscribe"), [{"T.AAPL", "T.NVDA"}])
+        await c.close()
+        await task
+
+    async def test_connect_again_subscribes_again(self):
+        # SLOW.X makes the server close the connection, so connect() returns
+        c = WebSocketClient(
+            api_key="", feed=f"localhost:{port}", secure=False, subscriptions=["SLOW.X"]
+        )
+        with self.assertLogs("WebSocketClient", level="WARNING"):
+            for _ in range(2):
+                await asyncio.wait_for(c.connect(lambda msgs: asyncio.sleep(0)), 5)
+        self.assertEqual(
+            [sent(conn, "subscribe") for conn in inbox], [[{"SLOW.X"}], [{"SLOW.X"}]]
+        )

test_websocket/test_status.py

@@ -2,30 +2,10 @@
 from massive.websocket import logger
 from base_ws import BaseTest
 from mock_server import port
-from websockets.asyncio.client import connect
 import asyncio
 
 
 class StatusTest(BaseTest):
-    async def asyncSetUp(self):
-        await super().asyncSetUp()
-        # wait until the mock server is accepting connections
-        for _ in range(100):
-            try:
-                async with connect(f"ws://localhost:{port}"):
-                    return
-            except OSError:
-                await asyncio.sleep(0.05)
-        self.fail("mock server did not start")
-
-    async def asyncTearDown(self):
-        # wait for the server to shut down so the next test can bind the port
-        self.task.cancel()
-        try:
-            await self.task
-        except asyncio.CancelledError:
-            pass
-
     async def run_client(self, *subscriptions):
         c = WebSocketClient(api_key="", feed=f"localhost:{port}", secure=False)
         c.subscribe(*subscriptions)