test: add end-to-end tests for the MQTT client

Author: allenporter

tests/conftest.py

@@ -1,4 +1,10 @@
+import io
+import logging
+import queue
 import re
+from collections.abc import Callable, Generator
+from typing import Any
+from unittest.mock import Mock, patch
 
 import pytest
 from aioresponses import aioresponses
@@ -8,9 +14,128 @@
 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
 from tests.mock_data import HOME_DATA_RAW, USER_DATA
 
+_LOGGER = logging.getLogger(__name__)
+
+
+# Used by fixtures to handle incoming requests and prepare responses
+RequestHandler = Callable[[bytes], bytes | None]
+
+
+class FakeSocketHandler:
+    """Fake socket used by the test to simulate a connection to the broker.
+
+    The socket handler is used to intercept the socket send and recv calls and
+    populate the response buffer with data to be sent back to the client. The
+    handle request callback handles the incoming requests and prepares the responses.
+    """
+
+    def __init__(self, handle_request: RequestHandler) -> None:
+        self.response_buf = io.BytesIO()
+        self.handle_request = handle_request
+
+    def pending(self) -> int:
+        """Return the number of bytes in the response buffer."""
+        return len(self.response_buf.getvalue())
+
+    def handle_socket_recv(self, read_size: int) -> bytes:
+        """Intercept a client recv() and populate the buffer."""
+        if self.pending() == 0:
+            raise BlockingIOError("No response queued")
+
+        self.response_buf.seek(0)
+        data = self.response_buf.read(read_size)
+        _LOGGER.debug("Response: 0x%s", data.hex())
+        # Consume the rest of the data in the buffer
+        remaining_data = self.response_buf.read()
+        self.response_buf = io.BytesIO(remaining_data)
+        return data
+
+    def handle_socket_send(self, client_request: bytes) -> int:
+        """Receive an incoming request from the client."""
+        _LOGGER.debug("Request: 0x%s", client_request.hex())
+        if (response := self.handle_request(client_request)) is not None:
+            # Enqueue a response to be sent back to the client in the buffer.
+            # The buffer will be emptied when the client calls recv() on the socket
+            _LOGGER.debug("Queued: 0x%s", response.hex())
+            self.response_buf.write(response)
+
+        return len(client_request)
+
+
+@pytest.fixture(name="received_requests")
+def received_requests_fixture() -> queue.Queue[bytes]:
+    """Fixture that provides access to the received requests."""
+    return queue.Queue()
+
+
+@pytest.fixture(name="response_queue")
+def response_queue_fixture() -> queue.Queue[bytes]:
+    """Fixture that provides access to the received requests."""
+    return queue.Queue()
+
+
+@pytest.fixture(name="request_handler")
+def request_handler_fixture(
+    received_requests: queue.Queue[bytes], response_queue: queue.Queue[bytes]
+) -> RequestHandler:
+    """Fixture records incoming requests and replies with responses from the queue."""
+
+    def handle_request(client_request: bytes) -> bytes | None:
+        """Handle an incoming request from the client."""
+        received_requests.put(client_request)
+
+        # Insert a prepared response into the response buffer
+        if response_queue.qsize() > 0:
+            return response_queue.get()
+        return None
+
+    return handle_request
+
+
+@pytest.fixture(name="fake_socket_handler")
+def fake_socket_handler_fixture(request_handler: RequestHandler) -> FakeSocketHandler:
+    """Fixture that creates a fake MQTT broker."""
+    return FakeSocketHandler(request_handler)
+
+
+@pytest.fixture(name="mock_sock")
+def mock_sock_fixture(fake_socket_handler: FakeSocketHandler) -> Mock:
+    """Fixture that creates a mock socket connection and wires it to the handler."""
+    mock_sock = Mock()
+    mock_sock.recv = fake_socket_handler.handle_socket_recv
+    mock_sock.send = fake_socket_handler.handle_socket_send
+    mock_sock.pending = fake_socket_handler.pending
+    return mock_sock
+
+
+@pytest.fixture(name="mock_create_connection")
+def create_connection_fixture(mock_sock: Mock) -> Generator[None, None, None]:
+    """Fixture that overrides the MQTT socket creation to wire it up to the mock socket."""
+    with patch("paho.mqtt.client.socket.create_connection", return_value=mock_sock):
+        yield
+
+
+@pytest.fixture(name="mock_select")
+def select_fixture(mock_sock: Mock, fake_socket_handler: FakeSocketHandler) -> Generator[None, None, None]:
+    """Fixture that overrides the MQTT client select calls to make select work on the mock socket.
+
+    This patch select to activate our mock socket when ready with data. Internal mqtt sockets are
+    always ready since they are used internally to wake the select loop. Ours is ready if there
+    is data in the buffer.
+    """
+
+    def is_ready(sock: Any) -> bool:
+        return sock is not mock_sock or (fake_socket_handler.pending() > 0)
+
+    def handle_select(rlist: list, wlist: list, *args: Any) -> list:
+        return [list(filter(is_ready, rlist)), list(filter(is_ready, wlist))]
+
+    with patch("paho.mqtt.client.select.select", side_effect=handle_select):
+        yield
+
 
 @pytest.fixture(name="mqtt_client")
-def mqtt_client():
+def mqtt_client(mock_create_connection: None, mock_select: None) -> Generator[RoborockMqttClientV1, None, None]:
     user_data = UserData.from_dict(USER_DATA)
     home_data = HomeData.from_dict(HOME_DATA_RAW)
     device_info = DeviceData(

tests/mock_data.py

@@ -21,14 +21,14 @@
         "r": {
             "r": "US",
             "a": "https://api-us.roborock.com",
-            "m": "ssl://mqtt-us.roborock.com:8883",
+            "m": "tcp://mqtt-us.roborock.com:8883",  # Skip SSL code in MQTT client library
             "l": "https://wood-us.roborock.com",
         },
     },
     "tuyaDeviceState": 2,
     "avatarurl": "https://files.roborock.com/iottest/default_avatar.png",
 }
-
+LOCAL_KEY = "key123"
 HOME_DATA_RAW = {
     "id": 123456,
     "name": "My Home",
@@ -199,7 +199,7 @@
             "name": "Roborock S7 MaxV",
             "attribute": None,
             "activeTime": 1672364449,
-            "localKey": "key123",
+            "localKey": LOCAL_KEY,
             "runtimeEnv": None,
             "timeZoneId": "America/Los_Angeles",
             "iconUrl": "no_url",
@@ -339,3 +339,5 @@
 }
 
 GET_CODE_RESPONSE = {"code": 200, "msg": "success", "data": None}
+
+MQTT_PUBLISH_TOPIC = "rr/m/o/user123/6ac2e6f8/abc123"

tests/mqtt_packet.py

@@ -0,0 +1,133 @@
+"""Module for crafting MQTT packets.
+
+This library is copied from the paho mqtt client library tests, with just the
+parts needed for some roborock messages. This message format in this file is
+not specific to roborock.
+"""
+
+import struct
+
+PROP_RECEIVE_MAXIMUM = 33
+PROP_TOPIC_ALIAS_MAXIMUM = 34
+
+
+def gen_uint16_prop(identifier: int, word: int) -> bytes:
+    """Generate a property with a uint16_t value."""
+    prop = struct.pack("!BH", identifier, word)
+    return prop
+
+
+def pack_varint(varint: int) -> bytes:
+    """Pack a variable integer."""
+    s = b""
+    while True:
+        byte = varint % 128
+        varint = varint // 128
+        # If there are more digits to encode, set the top bit of this digit
+        if varint > 0:
+            byte = byte | 0x80
+
+        s = s + struct.pack("!B", byte)
+        if varint == 0:
+            return s
+
+
+def prop_finalise(props: bytes) -> bytes:
+    """Finalise the properties."""
+    if props is None:
+        return pack_varint(0)
+    else:
+        return pack_varint(len(props)) + props
+
+
+def gen_connack(flags=0, rc=0, properties=b"", property_helper=True):
+    """Generate a CONNACK packet."""
+    if property_helper:
+        if properties is not None:
+            properties = (
+                gen_uint16_prop(PROP_TOPIC_ALIAS_MAXIMUM, 10) + properties + gen_uint16_prop(PROP_RECEIVE_MAXIMUM, 20)
+            )
+        else:
+            properties = b""
+    properties = prop_finalise(properties)
+
+    packet = struct.pack("!BBBB", 32, 2 + len(properties), flags, rc) + properties
+
+    return packet
+
+
+def gen_suback(mid: int, qos: int) -> bytes:
+    """Generate a SUBACK packet."""
+    return struct.pack("!BBHBB", 144, 2 + 1 + 1, mid, 0, qos)
+
+
+def _gen_short(cmd: int, reason_code: int) -> bytes:
+    return struct.pack("!BBB", cmd, 1, reason_code)
+
+
+def gen_disconnect(reason_code: int = 0) -> bytes:
+    """Generate a DISCONNECT packet."""
+    return _gen_short(0xE0, reason_code)
+
+
+def _gen_command_with_mid(cmd: int, mid: int, reason_code: int = 0) -> bytes:
+    return struct.pack("!BBHB", cmd, 3, mid, reason_code)
+
+
+def gen_puback(mid: int, reason_code: int = -1) -> bytes:
+    """Generate a PUBACK packet."""
+    return _gen_command_with_mid(64, mid, reason_code)
+
+
+def _pack_remaining_length(remaining_length: int) -> bytes:
+    """Pack a remaining length."""
+    s = b""
+    while True:
+        byte = remaining_length % 128
+        remaining_length = remaining_length // 128
+        # If there are more digits to encode, set the top bit of this digit
+        if remaining_length > 0:
+            byte = byte | 0x80
+
+        s = s + struct.pack("!B", byte)
+        if remaining_length == 0:
+            return s
+
+
+def gen_publish(
+    topic: str,
+    payload: bytes | None = None,
+    retain: bool = False,
+    dup: bool = False,
+    mid: int = 0,
+    properties: bytes = b"",
+) -> bytes:
+    """Generate a PUBLISH packet."""
+    if isinstance(topic, str):
+        topic_b = topic.encode("utf-8")
+    rl = 2 + len(topic_b)
+    pack_format = "H" + str(len(topic_b)) + "s"
+
+    properties = prop_finalise(properties)
+    rl += len(properties)
+    # This will break if len(properties) > 127
+    pack_format = pack_format + "%ds" % (len(properties))
+
+    if payload is not None:
+        # payload = payload.encode("utf-8")
+        rl = rl + len(payload)
+        pack_format = pack_format + str(len(payload)) + "s"
+    else:
+        payload = b""
+        pack_format = pack_format + "0s"
+
+    rlpacked = _pack_remaining_length(rl)
+    cmd = 48
+    if retain:
+        cmd = cmd + 1
+    if dup:
+        cmd = cmd + 8
+
+    return struct.pack(
+        "!B" + str(len(rlpacked)) + "s" + pack_format, cmd, rlpacked, len(topic_b), topic_b, properties, payload
+    )