chore: remove dependencies on get_request_id in RequestMessage

Author: allenporter

roborock/devices/v1_rpc_channel.py

@@ -14,12 +14,11 @@
 from roborock.protocols.v1_protocol import (
     CommandType,
     ParamsType,
+    RequestMessage,
     SecurityData,
-    create_mqtt_payload_encoder,
     decode_rpc_response,
-    encode_local_payload,
 )
-from roborock.roborock_message import RoborockMessage
+from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
 
 from .local_channel import LocalChannel
 from .mqtt_channel import MqttChannel
@@ -116,7 +115,7 @@ def __init__(
         self,
         name: str,
         channel: MqttChannel | LocalChannel,
-        payload_encoder: Callable[[CommandType, ParamsType], RoborockMessage],
+        payload_encoder: Callable[[RequestMessage], RoborockMessage],
     ) -> None:
         """Initialize the channel with a raw channel and an encoder function."""
         self._name = name
@@ -131,18 +130,26 @@ async def _send_raw_command(
     ) -> Any:
         """Send a command and return a parsed response RoborockBase type."""
         _LOGGER.debug("Sending command (%s): %s, params=%s", self._name, method, params)
-        message = self._payload_encoder(method, params)
+        request_message = RequestMessage(method, params=params)
+        message = self._payload_encoder(request_message)
         response = await self._channel.send_message(message)
         return decode_rpc_response(response)
 
 
 def create_mqtt_rpc_channel(mqtt_channel: MqttChannel, security_data: SecurityData) -> V1RpcChannel:
     """Create a V1 RPC channel using an MQTT channel."""
-    payload_encoder = create_mqtt_payload_encoder(security_data)
-    return PayloadEncodedV1RpcChannel("mqtt", mqtt_channel, payload_encoder)
+    return PayloadEncodedV1RpcChannel(
+        "mqtt",
+        mqtt_channel,
+        lambda x: x.encode_message(RoborockMessageProtocol.RPC_REQUEST, security_data=security_data),
+    )
 
 
 def create_combined_rpc_channel(local_channel: LocalChannel, mqtt_rpc_channel: V1RpcChannel) -> V1RpcChannel:
     """Create a V1 RPC channel that combines local and MQTT channels."""
-    local_rpc_channel = PayloadEncodedV1RpcChannel("local", local_channel, encode_local_payload)
+    local_rpc_channel = PayloadEncodedV1RpcChannel(
+        "local",
+        local_channel,
+        lambda x: x.encode_message(RoborockMessageProtocol.GENERAL_REQUEST),
+    )
     return CombinedV1RpcChannel(local_channel, local_rpc_channel, mqtt_rpc_channel)

roborock/protocols/v1_protocol.py

@@ -25,8 +25,6 @@
 __all__ = [
     "SecurityData",
     "create_security_data",
-    "create_mqtt_payload_encoder",
-    "encode_local_payload",
     "decode_rpc_response",
 ]
 
@@ -66,7 +64,19 @@ class RequestMessage:
     timestamp: int = field(default_factory=lambda: math.floor(time.time()))
     request_id: int = field(default_factory=lambda: get_next_int(10000, 32767))
 
-    def as_payload(self, security_data: SecurityData | None) -> bytes:
+    def encode_message(
+        self,
+        protocol: RoborockMessageProtocol,
+        security_data: SecurityData | None = None,
+    ) -> RoborockMessage:
+        """Convert the request message to a RoborockMessage."""
+        return RoborockMessage(
+            timestamp=self.timestamp,
+            protocol=protocol,
+            payload=self._as_payload(security_data=security_data),
+        )
+
+    def _as_payload(self, security_data: SecurityData | None) -> bytes:
         """Convert the request arguments to a dictionary."""
         inner = {
             "id": self.request_id,
@@ -85,35 +95,6 @@ def as_payload(self, security_data: SecurityData | None) -> bytes:
         )
 
 
-def create_mqtt_payload_encoder(security_data: SecurityData) -> Callable[[CommandType, ParamsType], RoborockMessage]:
-    """Create a payload encoder for V1 commands over MQTT."""
-
-    def _get_payload(method: CommandType, params: ParamsType) -> RoborockMessage:
-        """Build the payload for a V1 command."""
-        request = RequestMessage(method=method, params=params)
-        payload = request.as_payload(security_data)  # always secure
-        return RoborockMessage(
-            timestamp=request.timestamp,
-            protocol=RoborockMessageProtocol.RPC_REQUEST,
-            payload=payload,
-        )
-
-    return _get_payload
-
-
-def encode_local_payload(method: CommandType, params: ParamsType) -> RoborockMessage:
-    """Encode payload for V1 commands over local connection."""
-
-    request = RequestMessage(method=method, params=params)
-    payload = request.as_payload(security_data=None)
-
-    return RoborockMessage(
-        timestamp=request.timestamp,
-        protocol=RoborockMessageProtocol.GENERAL_REQUEST,
-        payload=payload,
-    )
-
-
 def decode_rpc_response(message: RoborockMessage) -> dict[str, Any]:
     """Decode a V1 RPC_RESPONSE message."""
     if not message.payload:

roborock/version_1_apis/roborock_local_client_v1.py

@@ -10,7 +10,7 @@
 from ..api import RoborockClient
 from ..exceptions import RoborockConnectionException, RoborockException, VacuumError
 from ..protocol import Decoder, Encoder, create_local_decoder, create_local_encoder
-from ..protocols.v1_protocol import encode_local_payload
+from ..protocols.v1_protocol import RequestMessage
 from ..roborock_message import RoborockMessage, RoborockMessageProtocol
 from ..util import RoborockLoggerAdapter
 from .roborock_client_v1 import CLOUD_REQUIRED, RoborockClientV1
@@ -123,12 +123,20 @@ async def async_disconnect(self) -> None:
 
     async def hello(self):
         try:
-            return await self._send_message(_HELLO_REQUEST_MESSAGE)
+            return await self._send_message(
+                roborock_message=_HELLO_REQUEST_MESSAGE,
+                request_id=_HELLO_REQUEST_MESSAGE.seq,
+                response_protocol=RoborockMessageProtocol.HELLO_RESPONSE,
+            )
         except Exception as e:
             self._logger.error(e)
 
     async def ping(self) -> None:
-        await self._send_message(_PING_REQUEST_MESSAGE)
+        await self._send_message(
+            roborock_message=_PING_REQUEST_MESSAGE,
+            request_id=_PING_REQUEST_MESSAGE.seq,
+            response_protocol=RoborockMessageProtocol.PING_RESPONSE,
+        )
 
     def _send_msg_raw(self, data: bytes):
         try:
@@ -145,27 +153,26 @@ async def _send_command(
     ):
         if method in CLOUD_REQUIRED:
             raise RoborockException(f"Method {method} is not supported over local connection")
-
-        roborock_message = encode_local_payload(method, params)
-        self._logger.debug("Building message id %s for method %s", roborock_message.get_request_id(), method)
-        return await self._send_message(roborock_message, method, params)
+        request_message = RequestMessage(method=method, params=params)
+        roborock_message = request_message.encode_message(RoborockMessageProtocol.GENERAL_REQUEST)
+        self._logger.debug("Building message id %s for method %s", request_message.request_id, method)
+        return await self._send_message(
+            roborock_message,
+            request_id=request_message.request_id,
+            response_protocol=RoborockMessageProtocol.GENERAL_REQUEST,
+            method=method,
+            params=params,
+        )
 
     async def _send_message(
         self,
         roborock_message: RoborockMessage,
+        request_id: int,
+        response_protocol: int,
         method: str | None = None,
         params: list | dict | int | None = None,
     ) -> RoborockMessage:
         await self.validate_connection()
-        request_id: int | None
-        if not method or not method.startswith("get"):
-            request_id = roborock_message.seq
-            response_protocol = request_id + 1
-        else:
-            request_id = roborock_message.get_request_id()
-            response_protocol = RoborockMessageProtocol.GENERAL_REQUEST
-        if request_id is None:
-            raise RoborockException(f"Failed build message {roborock_message}")
         msg = self._encoder(roborock_message)
         if method:
             self._logger.debug(f"id={request_id} Requesting method {method} with {params}")

roborock/version_1_apis/roborock_mqtt_client_v1.py

@@ -9,7 +9,7 @@
 
 from ..containers import DeviceData, UserData
 from ..exceptions import CommandVacuumError, RoborockException, VacuumError
-from ..protocols.v1_protocol import create_mqtt_payload_encoder, create_security_data
+from ..protocols.v1_protocol import RequestMessage, create_security_data
 from ..roborock_message import (
     RoborockMessageProtocol,
 )
@@ -28,12 +28,12 @@ def __init__(self, user_data: UserData, device_info: DeviceData, queue_timeout:
         rriot = user_data.rriot
         if rriot is None:
             raise RoborockException("Got no rriot data from user_data")
-        security_data = create_security_data(rriot)
         RoborockMqttClient.__init__(self, user_data, device_info)
+        security_data = create_security_data(rriot)
         RoborockClientV1.__init__(self, device_info, security_data=security_data)
         self.queue_timeout = queue_timeout
         self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER)
-        self._payload_encoder = create_mqtt_payload_encoder(security_data)
+        self._security_data = security_data
 
     async def _send_command(
         self,
@@ -44,13 +44,15 @@ async def _send_command(
             # When we have more custom commands do something more complicated here
             return await self._get_calibration_points()
 
-        roborock_message = self._payload_encoder(method, params)
+        request_message = RequestMessage(method=method, params=params)
+        roborock_message = request_message.encode_message(
+            RoborockMessageProtocol.RPC_REQUEST,
+            security_data=self._security_data,
+        )
         self._logger.debug("Building message id %s for method %s", roborock_message.get_request_id, method)
 
         await self.validate_connection()
-        request_id = roborock_message.get_request_id()
-        if request_id is None:
-            raise RoborockException(f"Failed build message {roborock_message}")
+        request_id = request_message.request_id
         response_protocol = (
             RoborockMessageProtocol.MAP_RESPONSE if method in COMMANDS_SECURED else RoborockMessageProtocol.RPC_RESPONSE
         )

tests/protocols/test_v1_protocol.py

@@ -10,11 +10,10 @@
 from roborock.exceptions import RoborockException
 from roborock.protocol import Utils
 from roborock.protocols.v1_protocol import (
+    RequestMessage,
     SecurityData,
     create_map_response_decoder,
-    create_mqtt_payload_encoder,
     decode_rpc_response,
-    encode_local_payload,
 )
 from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
 from roborock.roborock_typing import RoborockCommand
@@ -58,7 +57,7 @@ def request_id_fixture() -> Generator[int, None, None]:
 )
 def test_encode_local_payload(command, params, expected):
     """Test encoding of local payload for V1 commands."""
-    message = encode_local_payload(command, params)
+    message = RequestMessage(command, params).encode_message(RoborockMessageProtocol.GENERAL_REQUEST)
     assert isinstance(message, RoborockMessage)
     assert message.protocol == RoborockMessageProtocol.GENERAL_REQUEST
     assert message.payload == expected
@@ -76,8 +75,8 @@ def test_encode_local_payload(command, params, expected):
 )
 def test_encode_mqtt_payload(command, params, expected):
     """Test encoding of local payload for V1 commands."""
-    encoder = create_mqtt_payload_encoder(SECURITY_DATA)
-    message = encoder(command, params)
+    request_message = RequestMessage(command, params=params)
+    message = request_message.encode_message(RoborockMessageProtocol.RPC_REQUEST, SECURITY_DATA)
     assert isinstance(message, RoborockMessage)
     assert message.protocol == RoborockMessageProtocol.RPC_REQUEST
     assert message.payload == expected