chore: address some comments

Author: Lash-L

roborock/devices/b01_channel.py

@@ -33,8 +33,7 @@ async def send_decoded_command(
     _LOGGER.debug("Sending MQTT command: %s", params)
     msg_id = str(get_next_int(100000000000, 999999999999))
     roborock_message = encode_mqtt_payload(dps, command, params, msg_id)
-    finished = asyncio.Event()
-    result: dict[str, Any] = {}
+    future: asyncio.Future[dict[str, Any]] = asyncio.get_running_loop().create_future()
 
     def find_response(response_message: RoborockMessage) -> None:
         """Handle incoming messages and resolve the future."""
@@ -56,9 +55,11 @@ def find_response(response_message: RoborockMessage) -> None:
                 if isinstance(inner, dict) and inner.get("msgId") == msg_id:
                     _LOGGER.debug("Received query response: %s", inner)
                     data = inner.get("data")
-                    if isinstance(data, dict):
-                        result.update(data)
-                    finished.set()
+                    if not future.done():
+                        if isinstance(data, dict):
+                            future.set_result(data)
+                        else:
+                            future.set_exception(RoborockException(f"Unexpected data type for response: {data}"))
             else:
                 _LOGGER.debug("Received unexpected response: %s", dps_value)
 
@@ -68,10 +69,8 @@ def find_response(response_message: RoborockMessage) -> None:
     try:
         await mqtt_channel.publish(roborock_message)
         try:
-            await asyncio.wait_for(finished.wait(), timeout=_TIMEOUT)
+            return await asyncio.wait_for(future, timeout=_TIMEOUT)
         except TimeoutError as ex:
             raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex
     finally:
         unsub()
-
-    return result

roborock/protocols/b01_protocol.py

@@ -13,7 +13,6 @@
     RoborockMessage,
     RoborockMessageProtocol,
 )
-from roborock.util import get_next_int
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -28,7 +27,7 @@ def encode_mqtt_payload(dps: int, command: CommandType, params: ParamsType, msg_
         "dps": {
             dps: {
                 "method": str(command),
-                "msgId": msg_id or str(get_next_int(100000000000, 999999999999)),
+                "msgId": msg_id,
                 "params": params or [],
             }
         }

tests/devices/traits/a01/__init__.py

@@ -7,7 +7,9 @@
 from Crypto.Util.Padding import pad, unpad
 
 from roborock.data.b01_q7 import WorkStatusMapping
+from roborock.devices.b01_channel import send_decoded_command
 from roborock.devices.traits.b01.q7 import Q7PropertiesApi
+from roborock.exceptions import RoborockException
 from roborock.protocols.b01_protocol import B01_VERSION
 from roborock.roborock_message import RoborockB01Props, RoborockMessage, RoborockMessageProtocol
 from tests.conftest import FakeChannel
@@ -87,6 +89,7 @@ async def test_q7_api_query_values(q7_api: Q7PropertiesApi, fake_channel: FakeCh
     assert message.version == B01_VERSION
 
     # Verify request payload
+    assert message.payload is not None
     payload_data = json.loads(unpad(message.payload, AES.block_size))
     # {"dps": {"10000": {"method": "prop.get", "msgId": "123456789", "params": {"property": ["status", "wind"]}}}}
     assert "dps" in payload_data
@@ -128,3 +131,36 @@ async def test_q7_response_value_mapping(
         result = await q7_api.query_values(query)
 
     assert result is not None
+
+
+async def test_send_decoded_command_non_dict_response(fake_channel: FakeChannel):
+    """Test validity of handling non-dict responses (should not timeout)."""
+    msg_id = "123456789"
+
+    dps_payload = {
+        "dps": {
+            "10000": json.dumps(
+                {
+                    "msgId": msg_id,
+                    "data": "some_string_error",
+                }
+            )
+        }
+    }
+    message = RoborockMessage(
+        protocol=RoborockMessageProtocol.RPC_RESPONSE,
+        payload=pad(
+            json.dumps(dps_payload).encode(),
+            AES.block_size,
+        ),
+        version=b"B01",
+        seq=2021,
+    )
+
+    fake_channel.response_queue.append(message)
+
+    with patch("roborock.devices.b01_channel.get_next_int", return_value=int(msg_id)):
+        # Use a random string for command type to avoid needing import
+
+        with pytest.raises(RoborockException, match="Unexpected data type for response"):
+            await send_decoded_command(fake_channel, 10000, "prop.get", [])  # type: ignore[arg-type]