chore: add socket based tests for the new APIs (#677)

* chore: add socket based tests for the new APIs

This adds socket based tests for the new API so we can have tests that exercise the exact byte streams that are sent and received, similar to the old API. These currently are testing the lower level "session" and "channel" just to get the request infrastructure ready to be used in a real end to end device manaver tests.

This violates our rule that tests align with a module, but this is an exception given its for end to end tests.

* chore: Fix local session

Author: allenporter

tests/conftest.py

@@ -20,6 +20,11 @@
 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
 from tests.mock_data import HOME_DATA_RAW, HOME_DATA_SCENES_RAW, TEST_LOCAL_API_HOST, USER_DATA
 
+# Fixtures for the newer APIs in subdirectories
+pytest_plugins = [
+    "tests.mqtt_fixtures",
+]
+
 _LOGGER = logging.getLogger(__name__)
 
 

tests/e2e/__init__.py

@@ -0,0 +1 @@
+"""End-to-end tests package."""

tests/e2e/test_local_session.py

@@ -0,0 +1,138 @@
+"""End-to-end tests for LocalChannel using fake sockets."""
+
+import asyncio
+from collections.abc import AsyncGenerator, Callable, Generator
+from queue import Queue
+from typing import Any
+from unittest.mock import Mock, patch
+
+import pytest
+
+from roborock.devices.local_channel import LocalChannel
+from roborock.protocol import create_local_decoder, create_local_encoder
+from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
+from tests.conftest import RequestHandler
+from tests.mock_data import LOCAL_KEY
+
+TEST_HOST = "192.168.1.100"
+TEST_DEVICE_UID = "test_device_uid"
+TEST_CONNECT_NONCE = 12345
+TEST_ACK_NONCE = 67890
+TEST_RANDOM = 13579
+
+
+@pytest.fixture(name="mock_create_local_connection")
+def create_local_connection_fixture(request_handler: RequestHandler) -> Generator[None, None, None]:
+    """Fixture that overrides the transport creation to wire it up to the mock socket."""
+
+    async def create_connection(protocol_factory: Callable[[], asyncio.Protocol], *args, **kwargs) -> tuple[Any, Any]:
+        protocol = protocol_factory()
+
+        def handle_write(data: bytes) -> None:
+            response = request_handler(data)
+            if response is not None:
+                # Call data_received directly to avoid loop scheduling issues in test
+                protocol.data_received(response)
+
+        closed = asyncio.Event()
+
+        mock_transport = Mock()
+        mock_transport.write = handle_write
+        mock_transport.close = closed.set
+        mock_transport.is_reading = lambda: not closed.is_set()
+
+        return (mock_transport, protocol)
+
+    with patch("roborock.devices.local_channel.asyncio.get_running_loop") as mock_loop:
+        mock_loop.return_value.create_connection.side_effect = create_connection
+        yield
+
+
+@pytest.fixture(name="local_channel")
+async def local_channel_fixture(mock_create_local_connection: None) -> AsyncGenerator[LocalChannel, None]:
+    with patch(
+        "roborock.devices.local_channel.get_next_int", return_value=TEST_CONNECT_NONCE, device_uid=TEST_DEVICE_UID
+    ):
+        channel = LocalChannel(host=TEST_HOST, local_key=LOCAL_KEY, device_uid=TEST_DEVICE_UID)
+        yield channel
+        channel.close()
+
+
+def build_response(
+    protocol: RoborockMessageProtocol,
+    seq: int,
+    payload: bytes,
+    random: int,
+) -> bytes:
+    """Build an encoded response message."""
+    if protocol == RoborockMessageProtocol.HELLO_RESPONSE:
+        encoder = create_local_encoder(local_key=LOCAL_KEY, connect_nonce=TEST_CONNECT_NONCE, ack_nonce=None)
+    else:
+        encoder = create_local_encoder(local_key=LOCAL_KEY, connect_nonce=TEST_CONNECT_NONCE, ack_nonce=TEST_ACK_NONCE)
+
+    msg = RoborockMessage(
+        protocol=protocol,
+        random=random,
+        seq=seq,
+        payload=payload,
+    )
+    return encoder(msg)
+
+
+async def test_connect(
+    local_channel: LocalChannel, response_queue: Queue[bytes], received_requests: Queue[bytes]
+) -> None:
+    """Test connecting to the device."""
+    # Queue HELLO response with payload to ensure it can be parsed
+    response_queue.put(build_response(RoborockMessageProtocol.HELLO_RESPONSE, 1, payload=b"ok", random=TEST_RANDOM))
+
+    await local_channel.connect()
+
+    assert local_channel.is_connected
+    assert received_requests.qsize() == 1
+
+    # Verify HELLO request
+    request_bytes = received_requests.get()
+    # Note: We cannot use create_local_decoder here because HELLO_REQUEST has payload=None
+    # which causes MessageParser to fail parsing. For now we verify the raw bytes.
+
+    # Protocol is at offset 19 (2 bytes)
+    # Prefix(4) + Version(3) + Seq(4) + Random(4) + Timestamp(4) = 19
+    assert len(request_bytes) >= 21
+    protocol_bytes = request_bytes[19:21]
+    assert int.from_bytes(protocol_bytes, "big") == RoborockMessageProtocol.HELLO_REQUEST
+
+
+async def test_send_command(
+    local_channel: LocalChannel, response_queue: Queue[bytes], received_requests: Queue[bytes]
+) -> None:
+    """Test sending a command."""
+    # Queue HELLO response
+    response_queue.put(build_response(RoborockMessageProtocol.HELLO_RESPONSE, 1, payload=b"ok", random=TEST_RANDOM))
+
+    await local_channel.connect()
+
+    # Clear requests from handshake
+    while not received_requests.empty():
+        received_requests.get()
+
+    # Send command
+    cmd_seq = 123
+    msg = RoborockMessage(
+        protocol=RoborockMessageProtocol.RPC_REQUEST,
+        seq=cmd_seq,
+        payload=b'{"method":"get_status"}',
+    )
+
+    await local_channel.publish(msg)
+
+    # Verify request
+    assert received_requests.qsize() == 1
+    request_bytes = received_requests.get()
+
+    # Decode request
+    decoder = create_local_decoder(local_key=LOCAL_KEY, connect_nonce=TEST_CONNECT_NONCE, ack_nonce=TEST_ACK_NONCE)
+    msgs = list(decoder(request_bytes))
+    assert len(msgs) == 1
+    assert msgs[0].protocol == RoborockMessageProtocol.RPC_REQUEST
+    assert msgs[0].payload == b'{"method":"get_status"}'

tests/e2e/test_mqtt_session.py

@@ -0,0 +1,118 @@
+"""End-to-end tests for MQTT session.
+
+These tests use a fake MQTT broker to verify the session implementation. We
+mock out the lower level socket connections to simulate a broker which gets us
+close to an "end to end" test without needing an actual MQTT broker server.
+
+These are higher level tests that the similar tests in tests/mqtt/test_roborock_session.py
+which use mocks to verify specific behaviors.
+"""
+
+from collections.abc import AsyncGenerator, Callable
+from queue import Queue
+
+import pytest
+
+from roborock.mqtt.roborock_session import create_mqtt_session
+from roborock.mqtt.session import MqttSession
+from roborock.protocol import MessageParser
+from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
+from tests import mqtt_packet
+from tests.mock_data import LOCAL_KEY
+from tests.mqtt_fixtures import FAKE_PARAMS, Subscriber
+
+
+@pytest.fixture(autouse=True)
+def auto_mock_mqtt_client(mock_mqtt_client_fixture: None) -> None:
+    """Automatically use the mock mqtt client fixture."""
+
+
+@pytest.fixture(autouse=True)
+def auto_fast_backoff(fast_backoff_fixture: None) -> None:
+    """Automatically use the fast backoff fixture."""
+
+
+@pytest.fixture(autouse=True)
+def mqtt_server_fixture(mock_create_connection: None, mock_select: None) -> None:
+    """Fixture to mock the MQTT connection.
+
+    This is here to pull in the mock socket pixtures into all tests used here.
+    """
+
+
+@pytest.fixture(name="session")
+async def session_fixture(
+    push_response: Callable[[bytes], None],
+) -> AsyncGenerator[MqttSession, None]:
+    """Fixture to create a new connected MQTT session."""
+    push_response(mqtt_packet.gen_connack(rc=0, flags=2))
+    session = await create_mqtt_session(FAKE_PARAMS)
+    assert session.connected
+    try:
+        yield session
+    finally:
+        await session.close()
+
+
+async def test_session_e2e_receive_message(push_response: Callable[[bytes], None], session: MqttSession) -> None:
+    """Test receiving a real Roborock message through the session."""
+    assert session.connected
+
+    # Subscribe to the topic. We'll next construct and push a message.
+    push_response(mqtt_packet.gen_suback(mid=1))
+    subscriber = Subscriber()
+    await session.subscribe("topic-1", subscriber.append)
+
+    msg = RoborockMessage(
+        protocol=RoborockMessageProtocol.RPC_RESPONSE,
+        payload=b'{"result":"ok"}',
+        seq=123,
+    )
+    payload = MessageParser.build(msg, local_key=LOCAL_KEY, prefixed=False)
+
+    # Simulate receiving the message from the broker
+    push_response(mqtt_packet.gen_publish("topic-1", mid=2, payload=payload))
+
+    # Verify it was dispatched to the subscriber
+    await subscriber.wait()
+    assert len(subscriber.messages) == 1
+    received_payload = subscriber.messages[0]
+    assert received_payload == payload
+
+    # Verify the message payload contents
+    parsed_msgs, _ = MessageParser.parse(received_payload, local_key=LOCAL_KEY)
+    assert len(parsed_msgs) == 1
+    parsed_msg = parsed_msgs[0]
+    assert parsed_msg.protocol == RoborockMessageProtocol.RPC_RESPONSE
+    assert parsed_msg.seq == 123
+    # The payload in parsed_msg should be the decrypted bytes
+    assert parsed_msg.payload == b'{"result":"ok"}'
+
+
+async def test_session_e2e_publish_message(
+    push_response: Callable[[bytes], None],
+    received_requests: Queue,
+    session: MqttSession,
+) -> None:
+    """Test publishing a real Roborock message."""
+
+    # Publish a message to the brokwer
+    msg = RoborockMessage(
+        protocol=RoborockMessageProtocol.RPC_REQUEST,
+        payload=b'{"method":"get_status"}',
+        seq=456,
+    )
+    payload = MessageParser.build(msg, local_key=LOCAL_KEY, prefixed=False)
+
+    await session.publish("topic-1", payload)
+
+    # Verify what was sent to the broker
+    # We expect the payload to be present in the sent bytes
+    found = False
+    while not received_requests.empty():
+        request = received_requests.get()
+        if payload in request:
+            found = True
+            break
+
+    assert found, "Published payload not found in sent requests"