From 9b3b5595b7aa6687ce1c28e154777789ced1b549 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sat, 18 Jan 2025 17:31:22 -0800 Subject: [PATCH 1/3] test: add an additional local API test and fix bug in test fixture --- tests/conftest.py | 2 +- tests/test_local_api_v1.py | 65 +++++++++++++++++++++++++++++++++++--- 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 21bd5010..cf32c803 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -209,7 +209,7 @@ def handle_write(data: bytes) -> None: if response is not None: _LOGGER.debug("Replying with %s", response) loop = asyncio.get_running_loop() - loop.call_soon(protocol.data_received, response) + loop.call_soon(protocol.messages_cb, response) closed = asyncio.Event() diff --git a/tests/test_local_api_v1.py b/tests/test_local_api_v1.py index 5e22998d..b3069a45 100644 --- a/tests/test_local_api_v1.py +++ b/tests/test_local_api_v1.py @@ -1,21 +1,40 @@ """Tests for the Roborock Local Client V1.""" from queue import Queue +from unittest.mock import patch +import json +from typing import Any +from collections.abc import AsyncGenerator + +import pytest from roborock.protocol import MessageParser from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol from roborock.version_1_apis import RoborockLocalClientV1 +from roborock.containers import DeviceData, RoomMapping, S7MaxVStatus from .mock_data import LOCAL_KEY -def build_rpc_response(protocol: RoborockMessageProtocol, seq: int) -> bytes: +def build_rpc_response(seq: int, message: dict[str, Any]) -> bytes: + """Build an encoded RPC response message.""" + return build_raw_response( + protocol=RoborockMessageProtocol.GENERAL_REQUEST, + seq=seq, + payload=json.dumps( + { + "dps": {102: json.dumps(message)}, + } + ).encode(), + ) + +def build_raw_response(protocol: RoborockMessageProtocol, seq: int, payload: bytes) -> bytes: """Build an encoded RPC response message.""" message = RoborockMessage( protocol=protocol, random=23, seq=seq, - payload=b"ignored", + payload=payload, ) return MessageParser.build(message, local_key=LOCAL_KEY) @@ -26,8 +45,8 @@ async def test_async_connect( response_queue: Queue, ): """Test that we can connect to the Roborock device.""" - response_queue.put(build_rpc_response(RoborockMessageProtocol.HELLO_RESPONSE, 1)) - response_queue.put(build_rpc_response(RoborockMessageProtocol.PING_RESPONSE, 2)) + response_queue.put(build_raw_response(RoborockMessageProtocol.HELLO_RESPONSE, 1, b"ignored")) + response_queue.put(build_raw_response(RoborockMessageProtocol.PING_RESPONSE, 2, b"ignored")) await local_client.async_connect() assert local_client.is_connected() @@ -35,3 +54,41 @@ async def test_async_connect( await local_client.async_disconnect() assert not local_client.is_connected() + + + +@pytest.fixture(name="connected_local_client") +async def connected_local_client_fixture( + response_queue: Queue, local_client: RoborockLocalClientV1, +) -> AsyncGenerator[RoborockLocalClientV1, None]: + response_queue.put(build_raw_response(RoborockMessageProtocol.HELLO_RESPONSE, 1, b"ignored")) + response_queue.put(build_raw_response(RoborockMessageProtocol.PING_RESPONSE, 2, b"ignored")) + await local_client.async_connect() + yield local_client + + +async def test_get_room_mapping( + received_requests: Queue, + response_queue: Queue, + connected_local_client: RoborockLocalClientV1, +) -> None: + """Test sending an arbitrary MQTT message and parsing the response.""" + + test_request_id = 5050 + + message = build_rpc_response( + seq=test_request_id, + message={ + "id": test_request_id, + "result": [[16, "2362048"], [17, "2362044"]], + }, + ) + response_queue.put(message) + + with patch("roborock.version_1_apis.roborock_client_v1.get_next_int", return_value=test_request_id): + room_mapping = await connected_local_client.get_room_mapping() + + assert room_mapping == [ + RoomMapping(segment_id=16, iot_id="2362048"), + RoomMapping(segment_id=17, iot_id="2362044"), + ] From 35932b254cde12a18ab00da3cc32a9fc75b483df Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sat, 18 Jan 2025 17:35:00 -0800 Subject: [PATCH 2/3] test: fix formatting --- tests/test_local_api_v1.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/test_local_api_v1.py b/tests/test_local_api_v1.py index b3069a45..ab861d33 100644 --- a/tests/test_local_api_v1.py +++ b/tests/test_local_api_v1.py @@ -1,17 +1,17 @@ """Tests for the Roborock Local Client V1.""" -from queue import Queue -from unittest.mock import patch import json -from typing import Any from collections.abc import AsyncGenerator +from queue import Queue +from typing import Any +from unittest.mock import patch import pytest +from roborock.containers import RoomMapping from roborock.protocol import MessageParser from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol from roborock.version_1_apis import RoborockLocalClientV1 -from roborock.containers import DeviceData, RoomMapping, S7MaxVStatus from .mock_data import LOCAL_KEY @@ -28,6 +28,7 @@ def build_rpc_response(seq: int, message: dict[str, Any]) -> bytes: ).encode(), ) + def build_raw_response(protocol: RoborockMessageProtocol, seq: int, payload: bytes) -> bytes: """Build an encoded RPC response message.""" message = RoborockMessage( @@ -56,10 +57,10 @@ async def test_async_connect( assert not local_client.is_connected() - @pytest.fixture(name="connected_local_client") async def connected_local_client_fixture( - response_queue: Queue, local_client: RoborockLocalClientV1, + response_queue: Queue, + local_client: RoborockLocalClientV1, ) -> AsyncGenerator[RoborockLocalClientV1, None]: response_queue.put(build_raw_response(RoborockMessageProtocol.HELLO_RESPONSE, 1, b"ignored")) response_queue.put(build_raw_response(RoborockMessageProtocol.PING_RESPONSE, 2, b"ignored")) @@ -79,8 +80,8 @@ async def test_get_room_mapping( message = build_rpc_response( seq=test_request_id, message={ - "id": test_request_id, - "result": [[16, "2362048"], [17, "2362044"]], + "id": test_request_id, + "result": [[16, "2362048"], [17, "2362044"]], }, ) response_queue.put(message) From 52118fc965163fc27ef47f2c557e641ffba6a726 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sun, 19 Jan 2025 09:33:16 -0800 Subject: [PATCH 3/3] fix: Update local API --- roborock/local_api.py | 8 ++++++++ tests/conftest.py | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/roborock/local_api.py b/roborock/local_api.py index d3237056..b3064e0f 100644 --- a/roborock/local_api.py +++ b/roborock/local_api.py @@ -25,6 +25,14 @@ class _LocalProtocol(asyncio.Protocol): messages_cb: Callable[[bytes], None] connection_lost_cb: Callable[[Exception | None], None] + def data_received(self, bytes) -> None: + """Called when data is received from the transport.""" + self.messages_cb(bytes) + + def connection_lost(self, exc: Exception | None) -> None: + """Called when the transport connection is lost.""" + self.connection_lost_cb(exc) + class RoborockLocalClient(RoborockClient, ABC): """Roborock local client base class.""" diff --git a/tests/conftest.py b/tests/conftest.py index cf32c803..21bd5010 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -209,7 +209,7 @@ def handle_write(data: bytes) -> None: if response is not None: _LOGGER.debug("Replying with %s", response) loop = asyncio.get_running_loop() - loop.call_soon(protocol.messages_cb, response) + loop.call_soon(protocol.data_received, response) closed = asyncio.Event()