feat: Add exception handling and increased test coverage

Author: allenporter

roborock/mqtt/roborock_session.py

@@ -17,8 +17,7 @@
 import aiomqtt
 from aiomqtt import MqttError, TLSParameters
 
-from .. import RoborockException
-from .session import MqttParams, MqttSession
+from .session import MqttParams, MqttSession, MqttSessionException
 
 _LOGGER = logging.getLogger(__name__)
 _MQTT_LOGGER = logging.getLogger(f"{__name__}.aiomqtt")
@@ -71,7 +70,14 @@ async def start(self) -> None:
         start_future: asyncio.Future[None] = asyncio.Future()
         loop = asyncio.get_event_loop()
         self._background_task = loop.create_task(self._run_task(start_future))
-        await start_future
+        try:
+            await start_future
+        except MqttError as err:
+            raise MqttSessionException(f"Error starting MQTT session: {err}") from err
+        except Exception as err:
+            raise MqttSessionException(f"Unexpected error starting session: {err}") from err
+        else:
+            _LOGGER.debug("MQTT session started successfully")
 
     async def close(self) -> None:
         """Cancels the MQTT loop and shutdown the client library."""
@@ -102,14 +108,18 @@ async def _run_task(self, start_future: asyncio.Future[None] | None) -> None:
 
                     await self._process_message_loop(client)
 
-            except asyncio.CancelledError:
-                _LOGGER.debug("MQTT loop was cancelled")
-                return
             except MqttError as err:
-                _LOGGER.info("MQTT error: %s", err)
                 if start_future:
+                    _LOGGER.info("MQTT error starting session: %s", err)
                     start_future.set_exception(err)
                     return
+                _LOGGER.info("MQTT error: %s", err)
+            except asyncio.CancelledError as err:
+                if start_future:
+                    _LOGGER.debug("MQTT loop was cancelled")
+                    start_future.set_exception(err)
+                _LOGGER.debug("MQTT loop was cancelled whiel starting")
+                return
             # Catch exceptions to avoid crashing the loop
             # and to allow the loop to retry.
             except Exception as err:
@@ -118,10 +128,11 @@ async def _run_task(self, start_future: asyncio.Future[None] | None) -> None:
                 if "generator didn't stop" in str(err):
                     _LOGGER.debug("MQTT loop was cancelled")
                     return
-                _LOGGER.error("Uncaught error in MQTT session: %s", err)
                 if start_future:
+                    _LOGGER.error("Uncaught error starting MQTT session: %s", err)
                     start_future.set_exception(err)
                     return
+                _LOGGER.error("Uncaught error during MQTT session: %s", err)
 
             self._healthy = False
             _LOGGER.info("MQTT session disconnected, retrying in %s seconds", self._backoff.total_seconds())
@@ -150,6 +161,8 @@ async def _mqtt_client(self, params: MqttParams) -> aiomqtt.Client:
                     self._client = client
                     for topic in self._listeners:
                         _LOGGER.debug("Re-establising subscription to topic %s", topic)
+                        # TODO: If this fails it will break the whole connection. Make
+                        # this retry again in the background with backoff.
                         await client.subscribe(topic)
 
                 yield client
@@ -158,10 +171,11 @@ async def _mqtt_client(self, params: MqttParams) -> aiomqtt.Client:
                 self._client = None
 
     async def _process_message_loop(self, client: aiomqtt.Client) -> None:
-        _LOGGER.debug("Processing MQTT messages")
+        _LOGGER.debug("client=%s", client)
+        _LOGGER.debug("Processing MQTT messages: %s", client.messages)
         async for message in client.messages:
             _LOGGER.debug("Received message: %s", message)
-            for listener in self._listeners.get(message.topic.value) or []:
+            for listener in self._listeners.get(message.topic.value, []):
                 try:
                     listener(message.payload)
                 except asyncio.CancelledError:
@@ -185,7 +199,10 @@ async def subscribe(self, topic: str, callback: Callable[[bytes], None]) -> Call
         async with self._client_lock:
             if self._client:
                 _LOGGER.debug("Establishing subscription to topic %s", topic)
-                await self._client.subscribe(topic)
+                try:
+                    await self._client.subscribe(topic)
+                except MqttError as err:
+                    raise MqttSessionException(f"Error subscribing to topic: {err}") from err
             else:
                 _LOGGER.debug("Client not connected, will establish subscription later")
 
@@ -194,11 +211,15 @@ async def subscribe(self, topic: str, callback: Callable[[bytes], None]) -> Call
     async def publish(self, topic: str, message: bytes) -> None:
         """Publish a message on the topic."""
         _LOGGER.debug("Sending message to topic %s: %s", topic, message)
+        client: aiomqtt.Client
         async with self._client_lock:
-            if not self._client:
-                raise RoborockException("MQTT client not connected")
-            coro = self._client.publish(topic, message)
-        await coro
+            if self._client is None:
+                raise MqttSessionException("Could not publish message, MQTT client not connected")
+            client = self._client
+        try:
+            await client.publish(topic, message)
+        except MqttError as err:
+            raise MqttSessionException(f"Error publishing message: {err}") from err
 
 
 async def create_mqtt_session(params: MqttParams) -> MqttSession:

roborock/mqtt/session.py

@@ -4,6 +4,8 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 
+from roborock.exceptions import RoborockException
+
 DEFAULT_TIMEOUT = 30.0
 
 
@@ -55,3 +57,7 @@ async def publish(self, topic: str, message: bytes) -> None:
     @abstractmethod
     async def close(self) -> None:
         """Cancels the mqtt loop"""
+
+
+class MqttSessionException(RoborockException):
+    """ "Raised when there is an error communicating with MQTT."""

tests/mqtt/test_roborock_session.py

@@ -4,13 +4,14 @@
 from collections.abc import Callable, Generator
 from queue import Queue
 from typing import Any
-from unittest.mock import patch
+from unittest.mock import AsyncMock, Mock, patch
 
+import aiomqtt
 import paho.mqtt.client as mqtt
 import pytest
 
 from roborock.mqtt.roborock_session import create_mqtt_session
-from roborock.mqtt.session import MqttParams
+from roborock.mqtt.session import MqttParams, MqttSessionException
 from tests import mqtt_packet
 from tests.conftest import FakeSocketHandler
 
@@ -79,7 +80,11 @@ def push(message: bytes) -> None:
 
 
 class Subscriber:
-    """Mock subscriber class."""
+    """Mock subscriber class.
+
+    This will capture messages published on the session so the tests can verify
+    they were received.
+    """
 
     def __init__(self) -> None:
         """Initialize the subscriber."""
@@ -102,6 +107,7 @@ async def test_session(push_response: Callable[[bytes], None]) -> None:
 
     push_response(mqtt_packet.gen_connack(rc=0, flags=2))
     session = await create_mqtt_session(FAKE_PARAMS)
+    assert session.connected
 
     push_response(mqtt_packet.gen_suback(mid=1))
     subscriber1 = Subscriber()
@@ -130,7 +136,22 @@ async def test_session(push_response: Callable[[bytes], None]) -> None:
     push_response(mqtt_packet.gen_publish("topic-1", payload=b"ignored"))
     assert subscriber1.messages == [b"12345", b"ABC"]
 
+    assert session.connected
     await session.close()
+    assert not session.connected
+
+
+async def test_session_no_subscribers(push_response: Callable[[bytes], None]) -> None:
+    """Test the MQTT session."""
+
+    push_response(mqtt_packet.gen_connack(rc=0, flags=2))
+    push_response(mqtt_packet.gen_publish("topic-1", mid=3, payload=b"12345"))
+    push_response(mqtt_packet.gen_publish("topic-2", mid=4, payload=b"67890"))
+    session = await create_mqtt_session(FAKE_PARAMS)
+    assert session.connected
+
+    await session.close()
+    assert not session.connected
 
 
 async def test_publish_command(push_response: Callable[[bytes], None]) -> None:
@@ -139,4 +160,69 @@ async def test_publish_command(push_response: Callable[[bytes], None]) -> None:
     push_response(mqtt_packet.gen_connack(rc=0, flags=2))
     session = await create_mqtt_session(FAKE_PARAMS)
 
+    push_response(mqtt_packet.gen_publish("topic-1", mid=3, payload=b"12345"))
     await session.publish("topic-1", message=b"payload")
+
+    assert session.connected
+    await session.close()
+    assert not session.connected
+
+
+class FakeAsyncIterator:
+    """Fake async iterator that waits for messages to arrive, but they never do.
+
+    This is used for testing exceptions in other client functions.
+    """
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> None:
+        """Iterator that does not generate any messages."""
+        while True:
+            await asyncio.sleep(1)
+
+
+async def test_publish_failure() -> None:
+    """Test an MQTT error is received when publishing a message."""
+
+    mock_client = AsyncMock()
+    mock_client.messages = FakeAsyncIterator()
+
+    mock_aenter = AsyncMock()
+    mock_aenter.return_value = mock_client
+
+    with patch("roborock.mqtt.roborock_session.aiomqtt.Client.__aenter__", mock_aenter):
+        session = await create_mqtt_session(FAKE_PARAMS)
+        assert session.connected
+
+        mock_client.publish.side_effect = aiomqtt.MqttError
+
+        with pytest.raises(MqttSessionException, match="Error publishing message"):
+            await session.publish("topic-1", message=b"payload")
+
+
+async def test_subscribe_failure() -> None:
+    """Test an MQTT error while subscribing."""
+
+    mock_client = AsyncMock()
+    mock_client.messages = FakeAsyncIterator()
+
+    mock_aenter = AsyncMock()
+    mock_aenter.return_value = mock_client
+
+    mock_shim = Mock()
+    mock_shim.return_value.__aenter__ = mock_aenter
+    mock_shim.return_value.__aexit__ = AsyncMock()
+
+    with patch("roborock.mqtt.roborock_session.aiomqtt.Client", mock_shim):
+        session = await create_mqtt_session(FAKE_PARAMS)
+        assert session.connected
+
+        mock_client.subscribe.side_effect = aiomqtt.MqttError
+
+        subscriber1 = Subscriber()
+        with pytest.raises(MqttSessionException, match="Error subscribing to topic"):
+            await session.subscribe("topic-1", subscriber1.append)
+
+        assert not subscriber1.messages