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)