From 1f2d9879b7e653c09b3ecfa9004f0cf8705e3de7 Mon Sep 17 00:00:00 2001
From: Chip Kent <5250374+chipkent@users.noreply.github.com>
Date: Sat, 24 Oct 2026 00:06:35 -0600
Subject: [PATCH] fix: keep pending WebSocket subscriptions across a reconnect

subscribe() and unsubscribe() only change scheduled_subs, the topics
the caller wants. The receive loop sends the change on its next turn,
up to a second later, and then copies scheduled_subs into subs, the
topics the connection has been sent.

When the connection closed with an error, connect() set scheduled_subs
back to subs before reconnecting. That line came with "resub on
reconnect" (#207), and it does resubscribe the topics that were sent.
But a change made in the second before the drop had not been sent yet,
so it was thrown away: a topic subscribed just before the drop was
missing on the new connection, and a topic unsubscribed came back.

Keep scheduled_subs and only reset subs, so the new connection
subscribes to every topic the caller wants.

connect() also kept subs from an earlier call. Calling it again after
it returned, for example after the server closed the connection
cleanly, sent no subscribe message at all, because every topic looked
sent already. Reset subs at the start of every connect() call. This is
one way to end up with the missing subscriptions reported in #800.

The mock server now records the messages each connection receives, and
fails the connection on a DROP.X subscription. The shared test setup
waits for the server to start, as test_status did.
---
 massive/websocket/__init__.py    |  6 ++++-
 test_websocket/base_ws.py        | 28 +++++++++++++++++++-
 test_websocket/mock_server.py    | 14 +++++++++-
 test_websocket/test_reconnect.py | 45 ++++++++++++++++++++++++++++++++
 test_websocket/test_status.py    | 20 --------------
 5 files changed, 90 insertions(+), 23 deletions(-)
 create mode 100644 test_websocket/test_reconnect.py

diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py
index 7463b4cb..0f2d9610 100644
--- a/massive/websocket/__init__.py
+++ b/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:
diff --git a/test_websocket/base_ws.py b/test_websocket/base_ws.py
index a21bc4d8..54f74319 100644
--- a/test_websocket/base_ws.py
+++ b/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")
diff --git a/test_websocket/mock_server.py b/test_websocket/mock_server.py
index a3e4b294..db00e84a 100644
--- a/test_websocket/mock_server.py
+++ b/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(
diff --git a/test_websocket/test_reconnect.py b/test_websocket/test_reconnect.py
new file mode 100644
index 00000000..f7299c0a
--- /dev/null
+++ b/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"}]]
+        )
diff --git a/test_websocket/test_status.py b/test_websocket/test_status.py
index d333c9c8..f9ef275e 100644
--- a/test_websocket/test_status.py
+++ b/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)