fix: reconnect MQTT subscription with backoff after a dropped connection

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

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()