fix: minor changes

Author: humbertogontijo

roborock/api.py

@@ -23,12 +23,12 @@
     CleanSummary,
     Consumable,
     DNDTimer,
+    DeviceData,
     DustCollectionMode,
     HomeData,
     ModelStatus,
     MultiMapsList,
     NetworkInfo,
-    RoborockDeviceInfo,
     RoomMapping,
     SmartWashParams,
     Status,
@@ -86,7 +86,7 @@ async def request(self, method: str, url: str, params=None, data=None, headers=N
 
 
 class RoborockClient:
-    def __init__(self, endpoint: str, device_info: RoborockDeviceInfo) -> None:
+    def __init__(self, endpoint: str, device_info: DeviceData) -> None:
         self.device_info = device_info
         self._endpoint = endpoint
         self._nonce = secrets.token_bytes(16)
@@ -192,7 +192,12 @@ async def _async_response(self, request_id: int, protocol_id: int = 0) -> tuple[
         finally:
             del self._waiting_queue[request_id]
 
-    def _get_payload(self, method: RoborockCommand, params: Optional[list | dict] = None, secured=False):
+    def _get_payload(
+        self,
+        method: RoborockCommand,
+        params: Optional[list | dict] = None,
+        secured=False,
+    ):
         timestamp = math.floor(time.time())
         request_id = randint(10000, 99999)
         inner = {
@@ -244,7 +249,10 @@ async def get_clean_summary(self) -> CleanSummary | None:
             elif isinstance(clean_summary, list):
                 clean_time, clean_area, clean_count, records = unpack_list(clean_summary, 4)
                 return CleanSummary(
-                    clean_time=clean_time, clean_area=clean_area, clean_count=clean_count, records=records
+                    clean_time=clean_time,
+                    clean_area=clean_area,
+                    clean_count=clean_count,
+                    records=records,
                 )
         except RoborockTimeout as e:
             _LOGGER.error(e)

roborock/cli.py

@@ -10,7 +10,7 @@
 from roborock import RoborockException
 from roborock.api import RoborockApiClient
 from roborock.cloud_api import RoborockMqttClient
-from roborock.containers import LoginData, RoborockDeviceInfo
+from roborock.containers import LoginData, DeviceData
 from roborock.util import run_sync
 
 _LOGGER = logging.getLogger(__name__)
@@ -131,7 +131,7 @@ async def command(ctx, cmd, device_id, params):
     )
     if model is None:
         raise RoborockException(f"Could not find model for device {device.name}")
-    device_info = RoborockDeviceInfo(device=device, model=model)
+    device_info = DeviceData(device=device, model=model)
     mqtt_client = RoborockMqttClient(login_data.user_data, device_info)
     await mqtt_client.send_command(cmd, params)
     mqtt_client.__del__()

roborock/cloud_api.py

@@ -11,7 +11,7 @@
 import paho.mqtt.client as mqtt
 
 from .api import COMMANDS_SECURED, KEEPALIVE, RoborockClient, md5hex
-from .containers import RoborockDeviceInfo, UserData
+from .containers import DeviceData, UserData
 from .exceptions import CommandVacuumError, RoborockException, VacuumError
 from .protocol import MessageParser, Utils
 from .roborock_future import RoborockFuture
@@ -27,7 +27,7 @@ class RoborockMqttClient(RoborockClient, mqtt.Client):
     _thread: threading.Thread
     _client_id: str
 
-    def __init__(self, user_data: UserData, device_info: RoborockDeviceInfo) -> None:
+    def __init__(self, user_data: UserData, device_info: DeviceData) -> None:
         rriot = user_data.rriot
         if rriot is None:
             raise RoborockException("Got no rriot data from user_data")

roborock/containers.py

@@ -428,12 +428,19 @@ class NetworkInfo(RoborockBase):
 
 
 @dataclass
-class RoborockDeviceInfo(RoborockBase):
+class DeviceData(RoborockBase):
     device: HomeDataDevice
     model: str
+    host: Optional[str] = None
 
 
 @dataclass
 class RoomMapping(RoborockBase):
     segment_id: int
     iot_id: str
+
+
+@dataclass
+class BroadcastMessage(RoborockBase):
+    duid: str
+    ip: str

roborock/local_api.py

@@ -7,8 +7,8 @@
 
 import async_timeout
 
+from . import DeviceData
 from .api import COMMANDS_SECURED, QUEUE_TIMEOUT, RoborockClient
-from .containers import RoborockDeviceInfo
 from .exceptions import CommandVacuumError, RoborockConnectionException, RoborockException
 from .protocol import AP_CONFIG, MessageParser
 from .roborock_message import RoborockMessage
@@ -19,10 +19,12 @@
 
 
 class RoborockLocalClient(RoborockClient, asyncio.Protocol):
-    def __init__(self, device_info: RoborockDeviceInfo, ip: str):
-        super().__init__("abc", device_info)
+    def __init__(self, device_data: DeviceData):
+        if device_data.host is None:
+            raise RoborockException("Host is required")
+        super().__init__("abc", device_data)
         self.loop = get_running_loop_or_create_one()
-        self.ip = ip
+        self.host = device_data.host
         self._batch_structs: list[RoborockMessage] = []
         self._executing = False
         self.remaining = b""
@@ -47,13 +49,13 @@ async def async_connect(self) -> None:
             try:
                 if not self.is_connected():
                     async with async_timeout.timeout(QUEUE_TIMEOUT):
-                        _LOGGER.info(f"Connecting to {self.ip}")
+                        _LOGGER.info(f"Connecting to {self.host}")
                         self.transport, _ = await self.loop.create_connection(  # type: ignore
-                            lambda: self, self.ip, 58867
+                            lambda: self, self.host, 58867
                         )
-                        _LOGGER.info(f"Connected to {self.ip}")
+                        _LOGGER.info(f"Connected to {self.host}")
             except Exception as e:
-                raise RoborockConnectionException(f"Failed connecting to {self.ip}") from e
+                raise RoborockConnectionException(f"Failed connecting to {self.host}") from e
 
     def sync_disconnect(self) -> None:
         if self.transport and not self.loop.is_closed():

roborock/protocol.py

@@ -29,7 +29,7 @@
 from Crypto.Cipher import AES
 from Crypto.Util.Padding import pad, unpad
 
-from roborock import RoborockException
+from roborock import BroadcastMessage, RoborockException
 from roborock.roborock_message import RoborockMessage
 
 _LOGGER = logging.getLogger(__name__)
@@ -43,16 +43,16 @@ class RoborockProtocol(asyncio.DatagramProtocol):
     def __init__(self, timeout: int = 5):
         self.timeout = timeout
         self.transport: BaseTransport | None = None
-        self.devices_found: dict = {}
+        self.devices_found: list[BroadcastMessage] = []
         self._mutex = Lock()
 
     def __del__(self):
         self.close()
 
     def datagram_received(self, data, _):
         [broadcast_message], _ = BroadcastParser.parse(data)
-        parsed_message = json.loads(broadcast_message.payload)
-        self.devices_found[parsed_message.get("duid")] = parsed_message.get("ip")
+        parsed_message = BroadcastMessage.from_dict(json.loads(broadcast_message.payload))
+        self.devices_found.append(parsed_message)
 
     async def discover(self):
         async with self._mutex:
@@ -63,7 +63,7 @@ async def discover(self):
                 return self.devices_found
             finally:
                 self.close()
-                self.devices_found = {}
+                self.devices_found = []
 
     def close(self):
         self.transport.close() if self.transport else None

tests/conftest.py

@@ -2,15 +2,15 @@
 
 from roborock import HomeData, UserData
 from roborock.cloud_api import RoborockMqttClient
-from roborock.containers import RoborockDeviceInfo
+from roborock.containers import DeviceData
 from tests.mock_data import HOME_DATA_RAW, USER_DATA
 
 
 @pytest.fixture(name="mqtt_client")
 def mqtt_client():
     user_data = UserData.from_dict(USER_DATA)
     home_data = HomeData.from_dict(HOME_DATA_RAW)
-    device_info = RoborockDeviceInfo(
+    device_info = DeviceData(
         device=home_data.devices[0],
         model=home_data.products[0].model,
     )

tests/test_api.py

@@ -12,7 +12,7 @@
 )
 from roborock.api import PreparedRequest, RoborockApiClient
 from roborock.cloud_api import RoborockMqttClient
-from roborock.containers import RoborockDeviceInfo, S7MaxVStatus
+from roborock.containers import DeviceData, S7MaxVStatus
 from tests.mock_data import BASE_URL_REQUEST, GET_CODE_RESPONSE, HOME_DATA_RAW, STATUS, USER_DATA
 
 
@@ -26,7 +26,7 @@ def test_can_create_prepared_request():
 
 def test_can_create_mqtt_roborock():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
-    device_info = RoborockDeviceInfo(device=home_data.devices[0], model=home_data.products[0].model)
+    device_info = DeviceData(device=home_data.devices[0], model=home_data.products[0].model)
     RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
 
 
@@ -75,7 +75,7 @@ async def test_get_home_data():
 @pytest.mark.asyncio
 async def test_get_dust_collection_mode():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
-    device_info = RoborockDeviceInfo(device=home_data.devices[0], model=home_data.products[0].model)
+    device_info = DeviceData(device=home_data.devices[0], model=home_data.products[0].model)
     rmc = RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
     with patch("roborock.cloud_api.RoborockMqttClient.send_command") as command:
         command.return_value = {"mode": 1}
@@ -87,7 +87,7 @@ async def test_get_dust_collection_mode():
 @pytest.mark.asyncio
 async def test_get_mop_wash_mode():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
-    device_info = RoborockDeviceInfo(device=home_data.devices[0], model=home_data.products[0].model)
+    device_info = DeviceData(device=home_data.devices[0], model=home_data.products[0].model)
     rmc = RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
     with patch("roborock.cloud_api.RoborockMqttClient.send_command") as command:
         command.return_value = {"smart_wash": 0, "wash_interval": 1500}
@@ -100,7 +100,7 @@ async def test_get_mop_wash_mode():
 @pytest.mark.asyncio
 async def test_get_washing_mode():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
-    device_info = RoborockDeviceInfo(device=home_data.devices[0], model=home_data.products[0].model)
+    device_info = DeviceData(device=home_data.devices[0], model=home_data.products[0].model)
     rmc = RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
     with patch("roborock.cloud_api.RoborockMqttClient.send_command") as command:
         command.return_value = {"wash_mode": 2}
@@ -113,7 +113,7 @@ async def test_get_washing_mode():
 @pytest.mark.asyncio
 async def test_get_prop():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
-    device_info = RoborockDeviceInfo(device=home_data.devices[0], model=home_data.products[0].model)
+    device_info = DeviceData(device=home_data.devices[0], model=home_data.products[0].model)
     rmc = RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
     with patch("roborock.cloud_api.RoborockMqttClient.get_status") as get_status, patch(
         "roborock.cloud_api.RoborockMqttClient.send_command"