diff --git a/hyperliquid/websocket_manager.py b/hyperliquid/websocket_manager.py index 4c73a688..2225e324 100644 --- a/hyperliquid/websocket_manager.py +++ b/hyperliquid/websocket_manager.py @@ -127,7 +127,11 @@ def on_message(self, _ws, message): def on_open(self, _ws): logging.debug("on_open") self.ws_ready = True - for subscription, active_subscription in self.queued_subscriptions: + # Drain the queue before replaying it: subscribe() consults it for the + # single-subscription channels, and leaving entries behind would also + # replay them again on a later on_open. + queued_subscriptions, self.queued_subscriptions = self.queued_subscriptions, [] + for subscription, active_subscription in queued_subscriptions: self.subscribe(subscription, active_subscription.callback, active_subscription.subscription_id) def subscribe( @@ -136,16 +140,19 @@ def subscribe( if subscription_id is None: self.subscription_id_counter += 1 subscription_id = self.subscription_id_counter + identifier = subscription_to_identifier(subscription) + if identifier == "userEvents" or identifier == "orderUpdates": + # TODO: ideally the userEvent and orderUpdates messages would include the user so that we can multiplex + # Queued subscriptions count too, otherwise the duplicate is only caught while on_open replays the + # queue, where it aborts the replay and silently drops every subscription behind it. + already_queued = any(subscription_to_identifier(s) == identifier for s, _ in self.queued_subscriptions) + if len(self.active_subscriptions[identifier]) != 0 or already_queued: + raise NotImplementedError(f"Cannot subscribe to {identifier} multiple times") if not self.ws_ready: logging.debug("enqueueing subscription") self.queued_subscriptions.append((subscription, ActiveSubscription(callback, subscription_id))) else: logging.debug("subscribing") - identifier = subscription_to_identifier(subscription) - if identifier == "userEvents" or identifier == "orderUpdates": - # TODO: ideally the userEvent and orderUpdates messages would include the user so that we can multiplex - if len(self.active_subscriptions[identifier]) != 0: - raise NotImplementedError(f"Cannot subscribe to {identifier} multiple times") self.active_subscriptions[identifier].append(ActiveSubscription(callback, subscription_id)) self.ws.send(json.dumps({"method": "subscribe", "subscription": subscription})) return subscription_id diff --git a/tests/websocket_manager_test.py b/tests/websocket_manager_test.py new file mode 100644 index 00000000..6c4172d5 --- /dev/null +++ b/tests/websocket_manager_test.py @@ -0,0 +1,77 @@ +import json +from collections import defaultdict +from types import SimpleNamespace + +import pytest + +from hyperliquid.websocket_manager import WebsocketManager + + +def make_manager(ws_ready: bool): + """A WebsocketManager with a stub socket, so no connection is opened.""" + ws_manager = WebsocketManager.__new__(WebsocketManager) + ws_manager.subscription_id_counter = 0 + ws_manager.ws_ready = ws_ready + ws_manager.queued_subscriptions = [] + ws_manager.active_subscriptions = defaultdict(list) + sent = [] + ws_manager.ws = SimpleNamespace(send=sent.append) + return ws_manager, sent + + +def callback(_msg): + pass + + +def test_duplicate_single_subscription_raises_while_queued(): + # userEvents and orderUpdates cannot be multiplexed. Before this was checked on + # the queued path the duplicate was only caught later, inside on_open. + ws_manager, _ = make_manager(ws_ready=False) + ws_manager.subscribe({"type": "userEvents"}, callback) + + with pytest.raises(NotImplementedError): + ws_manager.subscribe({"type": "userEvents"}, callback) + + +def test_duplicate_single_subscription_raises_when_connected(): + ws_manager, _ = make_manager(ws_ready=True) + ws_manager.subscribe({"type": "orderUpdates"}, callback) + + with pytest.raises(NotImplementedError): + ws_manager.subscribe({"type": "orderUpdates"}, callback) + + +def test_rejected_duplicate_does_not_drop_later_subscriptions(): + # The duplicate used to surface inside on_open, which aborted the replay and + # silently dropped every subscription queued behind it. + ws_manager, sent = make_manager(ws_ready=False) + ws_manager.subscribe({"type": "userEvents"}, callback) + with pytest.raises(NotImplementedError): + ws_manager.subscribe({"type": "userEvents"}, callback) + ws_manager.subscribe({"type": "l2Book", "coin": "ETH"}, callback) + + ws_manager.on_open(None) + + assert len(ws_manager.active_subscriptions["l2Book:eth"]) == 1 + assert [json.loads(msg)["subscription"]["type"] for msg in sent] == ["userEvents", "l2Book"] + + +def test_on_open_replays_and_clears_the_queue(): + ws_manager, sent = make_manager(ws_ready=False) + ws_manager.subscribe({"type": "l2Book", "coin": "BTC"}, callback) + ws_manager.subscribe({"type": "trades", "coin": "ETH"}, callback) + + ws_manager.on_open(None) + + assert len(sent) == 2 + assert len(ws_manager.active_subscriptions["l2Book:btc"]) == 1 + assert len(ws_manager.active_subscriptions["trades:eth"]) == 1 + assert ws_manager.queued_subscriptions == [] + + +def test_multiplexable_channel_still_accepts_several_callbacks(): + ws_manager, _ = make_manager(ws_ready=True) + ws_manager.subscribe({"type": "l2Book", "coin": "ETH"}, callback) + ws_manager.subscribe({"type": "l2Book", "coin": "ETH"}, callback) + + assert len(ws_manager.active_subscriptions["l2Book:eth"]) == 2