fix: type checks

Author: humbertogontijo

roborock/api.py

@@ -122,7 +122,7 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
             self._last_device_msg_in = self.time_func()
             for data in messages:
                 protocol = data.protocol
-                if protocol == 102 or protocol == 4:
+                if data.payload and (protocol == 102 or protocol == 4):
                     payload = json.loads(data.payload.decode())
                     for data_point_number, data_point in payload.get("dps").items():
                         if data_point_number == "102":
@@ -146,7 +146,7 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
                                     if isinstance(result, list) and len(result) == 1:
                                         result = result[0]
                                     queue.resolve((result, None))
-                elif protocol == 301:
+                elif data.payload and protocol == 301:
                     payload = data.payload[0:24]
                     [endpoint, _, request_id, _] = struct.unpack("<15sBH6s", payload)
                     if endpoint.decode().startswith(self._endpoint):

roborock/local_api.py

@@ -98,13 +98,9 @@ async def ping(self):
         request_id = 2
         _LOGGER.debug(f"id={request_id} Requesting method ping with None")
         try:
-            return await self.send_message(RoborockMessage(
-                                protocol=2,
-                                payload=None,
-                                seq=request_id,
-                                version=b'1.0',
-                                random=23
-                                ))
+            return await self.send_message(
+                RoborockMessage(protocol=2, payload=None, seq=request_id, version=b"1.0", random=23)
+            )
         except Exception as e:
             _LOGGER.error(e)
 

roborock/protocol.py

@@ -53,9 +53,10 @@ def __del__(self):
 
     def datagram_received(self, data, _):
         [broadcast_message], _ = BroadcastParser.parse(data)
-        parsed_message = BroadcastMessage.from_dict(json.loads(broadcast_message.payload))
-        _LOGGER.debug(f"Received broadcast: {parsed_message}")
-        self.devices_found.append(parsed_message)
+        if broadcast_message.payload:
+            parsed_message = BroadcastMessage.from_dict(json.loads(broadcast_message.payload))
+            _LOGGER.debug(f"Received broadcast: {parsed_message}")
+            self.devices_found.append(parsed_message)
 
     async def discover(self):
         async with self._mutex:

roborock/roborock_message.py

@@ -6,8 +6,6 @@
 from dataclasses import dataclass
 from random import randint
 
-from .roborock_typing import RoborockCommand
-
 
 @dataclass
 class RoborockMessage:
@@ -20,17 +18,17 @@ class RoborockMessage:
 
     def get_request_id(self) -> int | None:
         protocol = self.protocol
-        if protocol in [4, 101, 102]:
+        if self.payload and protocol in [4, 101, 102]:
             payload = json.loads(self.payload.decode())
             for data_point_number, data_point in payload.get("dps").items():
                 if data_point_number in ["101", "102"]:
                     data_point_response = json.loads(data_point)
                     return data_point_response.get("id")
         return None
 
-    def get_method(self) -> RoborockCommand | None:
+    def get_method(self) -> str | None:
         protocol = self.protocol
-        if protocol in [4, 5, 101, 102]:
+        if self.payload and protocol in [4, 5, 101, 102]:
             payload = json.loads(self.payload.decode())
             for data_point_number, data_point in payload.get("dps").items():
                 if data_point_number in ["101", "102"]:
@@ -40,7 +38,7 @@ def get_method(self) -> RoborockCommand | None:
 
     def get_params(self) -> list | dict | None:
         protocol = self.protocol
-        if protocol in [4, 101, 102]:
+        if self.payload and protocol in [4, 101, 102]:
             payload = json.loads(self.payload.decode())
             for data_point_number, data_point in payload.get("dps").items():
                 if data_point_number in ["101", "102"]: