From 816661a3a92dca1c1f649fb51ab13070c8e0dcb3 Mon Sep 17 00:00:00 2001 From: Jasper Chan Date: Wed, 28 Oct 2026 14:51:48 -0700 Subject: [PATCH] fix: reconnect MQTT subscription with backoff after a dropped connection subscribe_mqtt ran a single aiomqtt session: any MqttError (timeout on connect, "Disconnected during message iteration") was logged and the task ended. Nothing re-established the subscription until the scheduled Cognito token refresh (~2.9 h later) or a manual config-entry reload, so Home Assistant silently stopped receiving activity_state pushes and every command failed after the 30 s wait in send_command. Fixes #47. - subscribe_mqtt now loops: each attempt runs one session via _run_mqtt_session and a failure is retried with exponential backoff (1 s .. 300 s), reset once a connection has stayed up for 60 s. - Each attempt reads the current id token, so a refreshed token is used on the next connect. - Cancellation (disconnect / schedule_reauthorization) still ends the task. - A failing callback or unparsable payload is logged instead of dropping the connection. - Log through the module logger instead of the root logger, so `python_snoo` debug logging shows the connection lifecycle. - Add pytest coverage for connect failure, mid-stream disconnect, _client_map bookkeeping, cancellation and token pick-up. Claude-Session: https://claude.ai/code/session_01Kjt2zjjJbbcYXsbXpzBUQk --- python_snoo/snoo.py | 76 ++++++++++++--- tests/__init__.py | 0 tests/test_mqtt_reconnect.py | 173 +++++++++++++++++++++++++++++++++++ 3 files changed, 234 insertions(+), 15 deletions(-) create mode 100644 tests/__init__.py create mode 100644 tests/test_mqtt_reconnect.py diff --git a/python_snoo/snoo.py b/python_snoo/snoo.py index f2895d5..69b36c8 100644 --- a/python_snoo/snoo.py +++ b/python_snoo/snoo.py @@ -19,6 +19,12 @@ _LOGGER = logging.getLogger(__name__) +# Reconnect policy for the MQTT subscription (see subscribe_mqtt). +MQTT_RECONNECT_MIN_DELAY = 1.0 +MQTT_RECONNECT_MAX_DELAY = 300.0 +# A connection that lived at least this long is considered healthy: the backoff resets after it drops. +MQTT_STABLE_CONNECTION_SECONDS = 60.0 + class Snoo: def __init__(self, email: str, password: str, clientsession: aiohttp.ClientSession): @@ -356,6 +362,47 @@ def start_subscribe(self, device: SnooDevice, function: Callable): self._mqtt_tasks[device.serialNumber] = asyncio.create_task(self.subscribe_mqtt(device, function)) async def subscribe_mqtt(self, device: SnooDevice, function: Callable): + """Keep an MQTT subscription for ``device`` alive until the task is cancelled. + + Every dropped or failed connection is retried with exponential backoff + (``MQTT_RECONNECT_MIN_DELAY`` .. ``MQTT_RECONNECT_MAX_DELAY`` seconds). The + backoff resets once a connection has stayed up for + ``MQTT_STABLE_CONNECTION_SECONDS``. Each attempt reads the current id token, + so a token refreshed by ``schedule_reauthorization`` is picked up automatically. + """ + delay = MQTT_RECONNECT_MIN_DELAY + while True: + connected_at: float | None = None + try: + connected_at = await self._run_mqtt_session(device, function) + except asyncio.CancelledError: + _LOGGER.debug(f"MQTT subscription task for {device.serialNumber} cancelled.") + raise + except aiomqtt.MqttError as e: + _LOGGER.warning(f"MQTT connection for {device.serialNumber} failed: {e}") + except Exception: + _LOGGER.exception(f"MQTT connection for {device.serialNumber} failed with an unexpected error") + finally: + # When the connection is lost, remove the client from the map so + # send_command waits for the next connection instead of using a dead client. + async with self._client_cond: + if device.serialNumber in self._client_map: + del self._client_map[device.serialNumber] + + if connected_at is not None: + uptime = asyncio.get_running_loop().time() - connected_at + if uptime >= MQTT_STABLE_CONNECTION_SECONDS: + delay = MQTT_RECONNECT_MIN_DELAY + _LOGGER.info(f"Reconnecting MQTT subscription for {device.serialNumber} in {delay:.0f} s.") + await asyncio.sleep(delay) + delay = min(delay * 2, MQTT_RECONNECT_MAX_DELAY) + + async def _run_mqtt_session(self, device: SnooDevice, function: Callable) -> float | None: + """Run one MQTT connection until it drops. + + Returns the loop time at which the connection was established, or ``None`` + if connecting failed. Exceptions propagate to ``subscribe_mqtt``. + """ host = device.awsIoT.clientEndpoint port = 443 websocket_path = "/mqtt" @@ -364,8 +411,9 @@ async def subscribe_mqtt(self, device: SnooDevice, function: Callable): password = None client_id = f"HA_{uuid.uuid4()}" user_name = "?SDK=iOS&Version=2.40.1" + connected_at: float | None = None - logging.debug(f"Attempting to connect to wss://{host}:{port}{websocket_path}") + _LOGGER.debug(f"Attempting to connect to wss://{host}:{port}{websocket_path}") # The default SSL context creation is a blocking I/O operation. # Run it in a separate thread to avoid blocking the Home Assistant event loop. @@ -385,7 +433,8 @@ async def subscribe_mqtt(self, device: SnooDevice, function: Callable): protocol=aiomqtt.ProtocolVersion.V31, timeout=10, ) as client: - logging.info(f"✅ Successfully connected to MQTT broker for {device.serialNumber}!") + connected_at = asyncio.get_running_loop().time() + _LOGGER.info(f"✅ Successfully connected to MQTT broker for {device.serialNumber}!") # Acquire the lock, add the client to the map, and notify waiting tasks. async with self._client_cond: @@ -394,19 +443,16 @@ async def subscribe_mqtt(self, device: SnooDevice, function: Callable): topic = f"{device.awsIoT.thingName}/state_machine/activity_state" await client.subscribe(topic) - logging.info(f"Subscribed to topic: {topic}") + _LOGGER.info(f"Subscribed to topic: {topic}") async for message in client.messages: - logging.debug(f"Received message on topic '{message.topic}': {message.payload.decode()}") - function(SnooData.from_json(message.payload.decode())) - - except aiomqtt.MqttError as e: - logging.error(f"MQTT connection for {device.serialNumber} failed: {e}") - except Exception as e: - logging.error(f"MQTT connection for {device.serialNumber} failed with an unexpected error: {e}") + payload = message.payload.decode() + _LOGGER.debug(f"Received message on topic '{message.topic}': {payload}") + try: + function(SnooData.from_json(payload)) + except Exception: + # A bad payload or a failing callback must not tear down the connection. + _LOGGER.exception(f"Error handling MQTT message for {device.serialNumber}") finally: - # When the connection is lost, remove the client from the map. - logging.info(f"MQTT connection closed for {device.serialNumber}.") - async with self._client_cond: - if device.serialNumber in self._client_map: - del self._client_map[device.serialNumber] + _LOGGER.info(f"MQTT connection closed for {device.serialNumber}.") + return connected_at diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_mqtt_reconnect.py b/tests/test_mqtt_reconnect.py new file mode 100644 index 0000000..3e1def7 --- /dev/null +++ b/tests/test_mqtt_reconnect.py @@ -0,0 +1,173 @@ +"""Tests for the MQTT subscription reconnect loop in python_snoo.snoo.Snoo.""" + +import asyncio +import json +from types import SimpleNamespace + +import aiomqtt +import pytest + +import python_snoo.snoo as snoo_module +from python_snoo.containers import AuthorizationInfo +from python_snoo.snoo import Snoo + +PAYLOAD = { + "system_state": "normal", + "sw_version": "v1.15.07", + "state_machine": { + "session_id": "1", + "state": "BASELINE", + "prev_state": "ONLINE", + "is_active_session": "true", + "since_session_start_ms": 1000, + "time_left": -1, + "hold": "off", + "weaning": "off", + "audio": "on", + "up_transition": "LEVEL1", + "down_transition": "NONE", + "sticky_white_noise": "off", + }, + "left_safety_clip": 1, + "right_safety_clip": 1, + "event": "command", + "event_time_ms": 1793223140196, + "rx_signal": {"rssi": -56, "strength": 78}, + "iot_capable": True, +} + + +class FakeClient: + """Stand-in for aiomqtt.Client driven by a script of connection outcomes. + + Each script entry is either ``"connect_error"`` (connecting raises MqttError) or an + int ``n`` (connect succeeds, ``n`` messages are delivered, then the connection + drops with MqttError, mirroring "Disconnected during message iteration"). + """ + + script: list = [] + instances: list = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.subscribed: list[str] = [] + self.outcome = FakeClient.script.pop(0) if FakeClient.script else "hang" + FakeClient.instances.append(self) + + async def __aenter__(self): + if self.outcome == "connect_error": + raise aiomqtt.MqttError("Operation timed out") + return self + + async def __aexit__(self, *exc): + return False + + async def subscribe(self, topic): + self.subscribed.append(topic) + + @property + def messages(self): + return self._messages() + + async def _messages(self): + if self.outcome == "hang": + await asyncio.Event().wait() # stay connected until cancelled + for _ in range(int(self.outcome)): + yield SimpleNamespace(topic="t/state_machine/activity_state", payload=json.dumps(PAYLOAD).encode()) + raise aiomqtt.MqttError("Disconnected during message iteration") + + +@pytest.fixture +def fast_backoff(monkeypatch): + monkeypatch.setattr(snoo_module, "MQTT_RECONNECT_MIN_DELAY", 0.01) + monkeypatch.setattr(snoo_module, "MQTT_RECONNECT_MAX_DELAY", 0.02) + monkeypatch.setattr(snoo_module, "MQTT_STABLE_CONNECTION_SECONDS", 0.0) + + +@pytest.fixture +def snoo(monkeypatch, fast_backoff): + FakeClient.script = [] + FakeClient.instances = [] + monkeypatch.setattr(snoo_module.aiomqtt, "Client", FakeClient) + s = Snoo("user@example.com", "pw", clientsession=None) + s.tokens = AuthorizationInfo(snoo="snoo-token", aws_access="a", aws_id="id-token", aws_refresh="r") + return s + + +def make_device(): + return SimpleNamespace( + serialNumber="123", + awsIoT=SimpleNamespace(clientEndpoint="example.invalid", thingName="thing_123"), + ) + + +async def wait_for(predicate, timeout=2.0): + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while not predicate(): + if loop.time() > deadline: + raise AssertionError("condition not met in time") + await asyncio.sleep(0.005) + + +@pytest.mark.asyncio +async def test_reconnects_after_connect_failure(snoo): + """A failed connect is retried instead of killing the subscription.""" + FakeClient.script = ["connect_error", 1, "hang"] + received = [] + snoo.start_subscribe(make_device(), received.append) + + await wait_for(lambda: len(received) == 1) + await wait_for(lambda: len(FakeClient.instances) == 3) + assert received[0].state_machine.state.value == "BASELINE" + assert FakeClient.instances[2].subscribed == ["thing_123/state_machine/activity_state"] + await snoo.disconnect() + + +@pytest.mark.asyncio +async def test_reconnects_after_mid_stream_disconnect(snoo): + """A connection that drops while iterating messages is re-established.""" + FakeClient.script = [1, 1, "hang"] + received = [] + snoo.start_subscribe(make_device(), received.append) + + await wait_for(lambda: len(received) == 2) + assert len(FakeClient.instances) >= 2 + await snoo.disconnect() + + +@pytest.mark.asyncio +async def test_client_map_reflects_connection_state(snoo): + """send_command must not see a dead client; the map is cleared on drop and refilled on reconnect.""" + FakeClient.script = [0, "hang"] + snoo.start_subscribe(make_device(), lambda data: None) + + await wait_for(lambda: len(FakeClient.instances) == 2) + await wait_for(lambda: snoo._client_map.get("123") is FakeClient.instances[1]) + await snoo.disconnect() + assert "123" not in snoo._client_map + + +@pytest.mark.asyncio +async def test_cancel_stops_reconnecting(snoo): + """Cancelling the task (disconnect / token refresh) ends the loop; no further connects.""" + FakeClient.script = ["connect_error", "connect_error", "connect_error"] + snoo.start_subscribe(make_device(), lambda data: None) + await wait_for(lambda: len(FakeClient.instances) >= 1) + await snoo.disconnect() + count = len(FakeClient.instances) + await asyncio.sleep(0.1) + assert len(FakeClient.instances) == count + assert snoo._mqtt_tasks == {} + + +@pytest.mark.asyncio +async def test_each_attempt_uses_current_token(snoo): + """A token refreshed between attempts is used by the next connection.""" + FakeClient.script = ["connect_error", "hang"] + snoo.start_subscribe(make_device(), lambda data: None) + await wait_for(lambda: len(FakeClient.instances) == 1) + snoo.tokens = AuthorizationInfo(snoo="snoo-token", aws_access="a", aws_id="new-id-token", aws_refresh="r") + await wait_for(lambda: len(FakeClient.instances) == 2) + assert FakeClient.instances[1].kwargs["websocket_headers"] == {"token": "new-id-token"} + await snoo.disconnect()