chore: Refactor to reuse the same payload functions

Author: allenporter

roborock/devices/v1_rpc_channel.py

@@ -21,6 +21,8 @@
     SecurityData,
     create_map_response_decoder,
     decode_rpc_response,
+    MapResponse,
+    ResponseMessage,
 )
 from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
 
@@ -114,155 +116,94 @@ async def _send_raw_command(
         raise RoborockException("No available connection to send command")
 
 
-class RpcPublisher:
-    """Helper to create send and receive messages on a channel."""
+class PayloadEncodedV1RpcChannel(BaseV1RpcChannel):
+    """Protocol for V1 channels that send encoded commands."""
 
     def __init__(
         self,
         name: str,
         channel: MqttChannel | LocalChannel,
         payload_encoder: Callable[[RequestMessage], RoborockMessage],
+        decoder: Callable[[RoborockMessage], ResponseMessage | MapResponse] | None = None,
     ) -> None:
-        """Initialize the RPC publisher."""
+        """Initialize the channel with a raw channel and an encoder function."""
         self._name = name
         self._channel = channel
         self._payload_encoder = payload_encoder
-
-    async def publish_and_wait(
-        self,
-        request_message: RequestMessage,
-        find_response: Callable[[RoborockMessage], None],
-        future: asyncio.Future[_V],
-    ) -> _V:
-        """Helper to send a message and wait for a future to complete.
-
-        The find_response function will be called for each incoming message. The
-        function should check the message and call future.set_result or
-        future.set_exception as appropriate when the response is found.
-        """
-        _LOGGER.debug(
-            "Sending command (%s, request_id=%s): %s, params=%s",
-            self._name,
-            request_message.request_id,
-            request_message.method,
-            request_message.params,
-        )
-        message = self._payload_encoder(request_message)
-        unsub = await self._channel.subscribe(find_response)
-        try:
-            await self._channel.publish(message)
-            return await asyncio.wait_for(future, timeout=_TIMEOUT)
-        except TimeoutError as ex:
-            future.cancel()
-            raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex
-        finally:
-            unsub()
-
-
-class PayloadEncodedV1RpcChannel(BaseV1RpcChannel):
-    """Protocol for V1 channels that send encoded commands."""
-
-    def __init__(self, publisher: RpcPublisher) -> None:
-        """Initialize the channel with a raw channel and an encoder function."""
-        self._name = publisher._name
-        self._publisher = publisher
+        self._decoder = decoder
 
     async def _send_raw_command(
         self,
         method: CommandType,
         *,
         params: ParamsType = None,
-    ) -> ResponseData:
+    ) -> ResponseData | bytes:
         """Send a command and return a parsed response RoborockBase type."""
         request_message = RequestMessage(method, params=params)
+        _LOGGER.debug(
+            "Sending command (%s, request_id=%s): %s, params=%s", self._name, request_message.request_id, method, params
+        )
+        message = self._payload_encoder(request_message)
 
-        future: asyncio.Future[ResponseData] = asyncio.Future()
+        future: asyncio.Future[ResponseData | bytes] = asyncio.Future()
 
         def find_response(response_message: RoborockMessage) -> None:
             try:
-                decoded = decode_rpc_response(response_message)
+                decoded = self._decoder(response_message)
             except RoborockException as ex:
                 _LOGGER.debug("Exception while decoding message (%s): %s", response_message, ex)
                 return
             _LOGGER.debug("Received response (%s, request_id=%s)", self._name, decoded.request_id)
             if decoded.request_id == request_message.request_id:
-                if decoded.api_error:
+                if isinstance(decoded, ResponseMessage) and decoded.api_error:
                     future.set_exception(decoded.api_error)
                 else:
                     future.set_result(decoded.data)
 
-        return await self._publisher.publish_and_wait(request_message, find_response, future)
-
-
-class MapRpcChannel(BaseV1RpcChannel):
-    """A V1 RPC channel that fetches and decodes map data."""
-
-    def __init__(
-        self,
-        publisher: RpcPublisher,
-        security_data: SecurityData,
-    ) -> None:
-        """Initialize the map RPC channel."""
-        self._publisher = publisher
-        self._decoder = create_map_response_decoder(security_data=security_data)
-
-    async def _send_raw_command(
-        self,
-        method: CommandType,
-        *,
-        params: ParamsType = None,
-    ) -> Any:
-        """Send a command and return a parsed response RoborockBase type."""
-        request_message = RequestMessage(method, params=params)
-
-        future: asyncio.Future[bytes] = asyncio.Future()
-
-        def find_response(response_message: RoborockMessage) -> None:
-            try:
-                decoded = self._decoder(response_message)
-            except RoborockException as ex:
-                _LOGGER.debug("Exception while decoding message (%s): %s", response_message, ex)
-                return
-            if decoded is None:
-                return
-            _LOGGER.debug("Received response (map), request_id=%s)", decoded.request_id)
-            if decoded.request_id == request_message.request_id:
-                future.set_result(decoded.data)
-
-        return await self._publisher.publish_and_wait(request_message, find_response, future)
-
+        message = self._payload_encoder(request_message)
+        unsub = await self._channel.subscribe(find_response)
+        try:
+            await self._channel.publish(message)
+            return await asyncio.wait_for(future, timeout=_TIMEOUT)
+        except TimeoutError as ex:
+            future.cancel()
+            raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex
+        finally:
+            unsub()
 
 def create_mqtt_rpc_channel(mqtt_channel: MqttChannel, security_data: SecurityData) -> V1RpcChannel:
     """Create a V1 RPC channel using an MQTT channel."""
-    publisher = RpcPublisher(
+    return PayloadEncodedV1RpcChannel(
         "mqtt",
         mqtt_channel,
         lambda x: x.encode_message(RoborockMessageProtocol.RPC_REQUEST, security_data=security_data),
+        decode_rpc_response,
     )
-    return PayloadEncodedV1RpcChannel(publisher)
 
 
 def create_local_rpc_channel(local_channel: LocalChannel) -> V1RpcChannel:
     """Create a V1 RPC channel using a local channel."""
-    publisher = RpcPublisher(
-        "local", local_channel, lambda x: x.encode_message(RoborockMessageProtocol.GENERAL_REQUEST)
+    return PayloadEncodedV1RpcChannel(
+        "local",
+        local_channel,
+        lambda x: x.encode_message(RoborockMessageProtocol.GENERAL_REQUEST),
+        decode_rpc_response,
     )
-    return PayloadEncodedV1RpcChannel(publisher)
 
 
 def create_map_rpc_channel(
     mqtt_channel: MqttChannel,
     security_data: SecurityData,
-) -> MapRpcChannel:
+) -> V1RpcChannel:
     """Create a V1 RPC channel that fetches map data.
 
     This will prefer local channels when available, falling back to MQTT
     channels if not. If neither is available, an exception will be raised
     when trying to send a command.
     """
-    publisher = RpcPublisher(
+    return PayloadEncodedV1RpcChannel(
         "map",
         mqtt_channel,
         lambda x: x.encode_message(RoborockMessageProtocol.RPC_REQUEST, security_data=security_data),
+        create_map_response_decoder(security_data=security_data),
     )
-    return MapRpcChannel(publisher, security_data)