feat: Add test coverage to device modules

Author: allenporter

roborock/devices/device.py

@@ -6,6 +6,7 @@
 
 import enum
 import logging
+from collections.abc import Callable
 from functools import cached_property
 
 from roborock.containers import HomeDataDevice, HomeDataProduct, UserData
@@ -40,13 +41,16 @@ def __init__(
     ) -> None:
         """Initialize the RoborockDevice.
 
-        The device takes ownership of the MQTT channel for communication with the device
-        and will close it when the device is closed.
+        The device takes ownership of the MQTT channel for communication with the device.
+        Use `connect()` to establish the connection, which will set up the MQTT channel
+        for receiving messages from the device. Use `close()` to unsubscribe from the MQTT
+        channel.
         """
         self._user_data = user_data
         self._device_info = device_info
         self._product_info = product_info
         self._mqtt_channel = mqtt_channel
+        self._unsub: Callable[[], None] | None = None
 
     @property
     def duid(self) -> str:
@@ -82,14 +86,18 @@ async def connect(self) -> None:
 
         This method will set up the MQTT channel for communication with the device.
         """
-        await self._mqtt_channel.subscribe(self._on_mqtt_message)
+        if self._unsub:
+            raise ValueError("Already connected to the device")
+        self._unsub = await self._mqtt_channel.subscribe(self._on_mqtt_message)
 
     async def close(self) -> None:
         """Close the MQTT connection to the device.
 
         This method will unsubscribe from the MQTT channel and clean up resources.
         """
-        await self._mqtt_channel.close()
+        if self._unsub:
+            self._unsub()
+            self._unsub = None
 
     def _on_mqtt_message(self, message: bytes) -> None:
         """Handle incoming MQTT messages from the device.

roborock/devices/mqtt_channel.py

@@ -31,14 +31,14 @@ def _subscribe_topic(self) -> str:
         """Topic to receive responses from the device."""
         return f"rr/m/o/{self._rriot.u}/{self._mqtt_params.username}/{self._duid}"
 
-    async def subscribe(self, callback: Callable[[bytes], None]) -> None:
-        """Subscribe to the device's response topic."""
-        if self._unsub:
-            raise ValueError("Already subscribed to the response topic")
-        self._unsub = await self._mqtt_session.subscribe(self._subscribe_topic, callback)
+    async def subscribe(self, callback: Callable[[bytes], None]) -> Callable[[], None]:
+        """Subscribe to the device's response topic.
 
-    async def close(self) -> None:
-        """Close the MQTT subscription."""
+        The callback will be called with the message payload when a message is received.
+        If already subscribed, raises ValueError.
+
+        Returns a callable that can be used to unsubscribe from the topic.
+        """
         if self._unsub:
-            self._unsub()
-            self._unsub = None
+            raise ValueError("Already subscribed to the response topic")
+        return await self._mqtt_session.subscribe(self._subscribe_topic, callback)

tests/devices/test_device.py

@@ -0,0 +1,40 @@
+"""Tests for the Device class."""
+
+from unittest.mock import AsyncMock, Mock
+
+from roborock.containers import HomeData, UserData
+from roborock.devices.device import DeviceVersion, RoborockDevice
+
+from .. import mock_data
+
+USER_DATA = UserData.from_dict(mock_data.USER_DATA)
+HOME_DATA = HomeData.from_dict(mock_data.HOME_DATA_RAW)
+
+
+async def test_device_connection() -> None:
+    """Test the Device connection setup."""
+
+    unsub = Mock()
+    subscribe = AsyncMock()
+    subscribe.return_value = unsub
+    mqtt_channel = AsyncMock()
+    mqtt_channel.subscribe = subscribe
+
+    device = RoborockDevice(
+        USER_DATA,
+        device_info=HOME_DATA.devices[0],
+        product_info=HOME_DATA.products[0],
+        mqtt_channel=mqtt_channel,
+    )
+    assert device.duid == "abc123"
+    assert device.name == "Roborock S7 MaxV"
+    assert device.device_version == DeviceVersion.V1
+
+    assert not subscribe.called
+
+    await device.connect()
+    assert subscribe.called
+    assert not unsub.called
+
+    await device.close()
+    assert unsub.called

tests/devices/test_device_manager.py

@@ -1,5 +1,6 @@
 """Tests for the DeviceManager class."""
 
+from collections.abc import Generator
 from unittest.mock import patch
 
 import pytest
@@ -14,6 +15,13 @@
 USER_DATA = UserData.from_dict(mock_data.USER_DATA)
 
 
+@pytest.fixture(autouse=True)
+def setup_mqtt_session() -> Generator[None, None, None]:
+    """Fixture to set up the MQTT session for the tests."""
+    with patch("roborock.devices.device_manager.create_mqtt_session"):
+        yield
+
+
 async def home_home_data_no_devices() -> HomeData:
     """Mock home data API that returns no devices."""
     return HomeData(

tests/devices/test_mqtt_channel.py

@@ -0,0 +1,64 @@
+"""Tests for the MqttChannel class."""
+
+from collections.abc import Generator
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from roborock.containers import HomeData, UserData
+from roborock.devices.mqtt_channel import MqttChannel
+from roborock.mqtt.session import MqttParams
+
+from .. import mock_data
+
+USER_DATA = UserData.from_dict(mock_data.USER_DATA)
+TEST_MQTT_PARAMS = MqttParams(
+    host="localhost",
+    port=1883,
+    tls=False,
+    username="username",
+    password="password",
+    timeout=10.0,
+)
+
+
+@pytest.fixture(autouse=True)
+def setup_mqtt_session() -> Generator[None, None, None]:
+    """Fixture to set up the MQTT session for the tests."""
+    with patch("roborock.devices.device_manager.create_mqtt_session"):
+        yield
+
+
+async def home_home_data_no_devices() -> HomeData:
+    """Mock home data API that returns no devices."""
+    return HomeData(
+        id=1,
+        name="Test Home",
+        devices=[],
+        products=[],
+    )
+
+
+async def mock_home_data() -> HomeData:
+    """Mock home data API that returns devices."""
+    return HomeData.from_dict(mock_data.HOME_DATA_RAW)
+
+
+async def test_mqtt_channel() -> None:
+    """Test MQTT channel setup."""
+
+    mock_session = AsyncMock()
+
+    channel = MqttChannel(mock_session, duid="abc123", rriot=USER_DATA.rriot, mqtt_params=TEST_MQTT_PARAMS)
+
+    unsub = Mock()
+    mock_session.subscribe.return_value = unsub
+
+    callback = Mock()
+    result = await channel.subscribe(callback)
+
+    assert mock_session.subscribe.called
+    assert mock_session.subscribe.call_args[0][0] == "rr/m/o/user123/username/abc123"
+    assert mock_session.subscribe.call_args[0][1] == callback
+
+    assert result == unsub