From 862be93f9298fc746a2fd710c098235195dde42d Mon Sep 17 00:00:00 2001 From: Paul Bottein Date: Fri, 24 Jul 2026 22:11:42 +0200 Subject: [PATCH 1/2] fix(a01): restart stale MQTT session on command timeouts --- roborock/devices/rpc/a01_channel.py | 2 ++ roborock/devices/traits/a01/__init__.py | 40 +++++++++++++++++++++++-- tests/devices/rpc/test_a01_channel.py | 36 ++++++++++++++++++++++ tests/devices/traits/a01/test_init.py | 21 +++++++++++++ 4 files changed, 97 insertions(+), 2 deletions(-) diff --git a/roborock/devices/rpc/a01_channel.py b/roborock/devices/rpc/a01_channel.py index 66b8b27d7..4293df5f5 100644 --- a/roborock/devices/rpc/a01_channel.py +++ b/roborock/devices/rpc/a01_channel.py @@ -98,8 +98,10 @@ def find_response(response_message: RoborockMessage) -> None: try: await asyncio.wait_for(finished.wait(), timeout=_TIMEOUT) except TimeoutError as ex: + await mqtt_channel.health_manager.on_timeout() raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex finally: unsub() + await mqtt_channel.health_manager.on_success() return result # type: ignore[return-value] diff --git a/roborock/devices/traits/a01/__init__.py b/roborock/devices/traits/a01/__init__.py index 537c84377..739ef30f3 100644 --- a/roborock/devices/traits/a01/__init__.py +++ b/roborock/devices/traits/a01/__init__.py @@ -8,12 +8,14 @@ A01 devices expose a single API object that handles all device interactions. This API is available on the device instance (typically via `device.a01_properties`). -The API provides two main methods: +The API provides these methods: 1. **query_values(protocols)**: Fetches current state for specific data points. You must pass a list of protocol enums (e.g. `RoborockDyadDataProtocol` or `RoborockZeoProtocol`) to request specific data. 2. **set_value(protocol, value)**: Sends a command to the device to change a setting or perform an action. +3. **add_listener(callback)**: Subscribes to state the device pushes on its own (for + example when its state changes), invoking the callback with decoded values. Note that these APIs fetch data directly from the device upon request and do not cache state internally. @@ -51,7 +53,9 @@ from roborock.devices.rpc.a01_channel import send_decoded_command from roborock.devices.traits import Trait from roborock.devices.transport.mqtt_channel import MqttChannel -from roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol +from roborock.exceptions import RoborockException +from roborock.protocols.a01_protocol import decode_rpc_response +from roborock.roborock_message import RoborockDyadDataProtocol, RoborockMessage, RoborockZeoProtocol __init__ = [ "DyadApi", @@ -134,6 +138,12 @@ def convert_zeo_value(protocol_value: RoborockZeoProtocol, value: Any) -> Any: return None +# RoborockDyadDataProtocol._missing_ maps any unknown code to its first member +# instead of raising, so incoming data points must be checked against this set +# before being converted to a protocol. +_DYAD_PROTOCOL_VALUES = frozenset(protocol.value for protocol in RoborockDyadDataProtocol) + + class DyadApi(Trait): """API for interacting with Dyad devices.""" @@ -155,6 +165,32 @@ async def set_value(self, protocol: RoborockDyadDataProtocol, value: Any) -> dic params = {protocol: value} return await send_decoded_command(self._channel, params) + async def add_listener( + self, callback: Callable[[dict[RoborockDyadDataProtocol, Any]], None] + ) -> Callable[[], None]: + """Listen for state the device pushes on its own. + + The callback is invoked with decoded values whenever the device sends a + message, including unsolicited pushes when its state changes. Only known + protocols are delivered. Returns a callable to remove the listener. + """ + + def on_message(message: RoborockMessage) -> None: + try: + datapoints = decode_rpc_response(message) + except RoborockException: + return + values: dict[RoborockDyadDataProtocol, Any] = {} + for code, value in datapoints.items(): + if code not in _DYAD_PROTOCOL_VALUES: + continue + protocol = RoborockDyadDataProtocol(code) + values[protocol] = convert_dyad_value(protocol, value) + if values: + callback(values) + + return await self._channel.subscribe(on_message) + class ZeoApi(Trait): """API for interacting with Zeo devices.""" diff --git a/tests/devices/rpc/test_a01_channel.py b/tests/devices/rpc/test_a01_channel.py index e98cf228e..9f2f5a767 100644 --- a/tests/devices/rpc/test_a01_channel.py +++ b/tests/devices/rpc/test_a01_channel.py @@ -1,10 +1,12 @@ """Tests for the a01_channel.""" from typing import Any +from unittest.mock import AsyncMock import pytest from roborock.devices.rpc.a01_channel import send_decoded_command +from roborock.exceptions import RoborockException from roborock.protocols.a01_protocol import encode_mqtt_payload from roborock.roborock_message import ( RoborockDyadDataProtocol, @@ -51,3 +53,37 @@ async def test_id_query(mock_mqtt_channel: FakeChannel): } mock_mqtt_channel.publish.assert_awaited_once() mock_mqtt_channel.subscribe.assert_awaited_once() + + +async def test_query_marks_session_healthy(mock_mqtt_channel: FakeChannel): + """A completed query reports success to the health manager.""" + mock_mqtt_channel.health_manager.on_success = AsyncMock() # type: ignore[method-assign] + + params: dict[RoborockDyadDataProtocol, Any] = { + RoborockDyadDataProtocol.ID_QUERY: [RoborockDyadDataProtocol.POWER] + } + encoded = encode_mqtt_payload({RoborockDyadDataProtocol.POWER: 75}, value_encoder=lambda x: x) + mock_mqtt_channel.response_queue.append( + RoborockMessage( + protocol=RoborockMessageProtocol.RPC_RESPONSE, payload=encoded.payload, version=encoded.version + ) + ) + + await send_decoded_command(mock_mqtt_channel, params) # type: ignore[call-overload] + + mock_mqtt_channel.health_manager.on_success.assert_awaited_once() + + +async def test_query_timeout_reports_to_health_manager(mock_mqtt_channel: FakeChannel, monkeypatch): + """A timed-out query reports to the health manager so a stale session can be restarted.""" + monkeypatch.setattr("roborock.devices.rpc.a01_channel._TIMEOUT", 0.01) + mock_mqtt_channel.health_manager.on_timeout = AsyncMock() # type: ignore[method-assign] + + params: dict[RoborockDyadDataProtocol, Any] = { + RoborockDyadDataProtocol.ID_QUERY: [RoborockDyadDataProtocol.POWER] + } + + with pytest.raises(RoborockException, match="timed out"): + await send_decoded_command(mock_mqtt_channel, params) # type: ignore[call-overload] + + mock_mqtt_channel.health_manager.on_timeout.assert_awaited_once() diff --git a/tests/devices/traits/a01/test_init.py b/tests/devices/traits/a01/test_init.py index 7e2fd6a69..ecdb92474 100644 --- a/tests/devices/traits/a01/test_init.py +++ b/tests/devices/traits/a01/test_init.py @@ -136,6 +136,27 @@ async def test_dyad_invalid_response_value( assert result == expected_result +async def test_dyad_add_listener(dyad_api: DyadApi, fake_channel: FakeChannel): + """add_listener delivers decoded values for pushed messages and skips unknown codes.""" + received: list[dict[RoborockDyadDataProtocol, Any]] = [] + unsub = await dyad_api.add_listener(received.append) + + # 999 is not a known protocol: it must be skipped, not mapped to the first enum member. + fake_channel.notify_subscribers(build_a01_message({206: 3, 209: 80, 216: 0, 999: 1})) + + assert received == [ + { + RoborockDyadDataProtocol.SUCTION: "l3", + RoborockDyadDataProtocol.POWER: 80, + RoborockDyadDataProtocol.ERROR: "none", + } + ] + + unsub() + fake_channel.notify_subscribers(build_a01_message({206: 1})) + assert len(received) == 1 # no callback fires after unsubscribing + + async def test_zeo_api_query_values(zeo_api: ZeoApi, fake_channel: FakeChannel): """Test that ZeoApi currently returns raw values without conversion.""" fake_channel.response_queue.append( From 7b26533f360c154d679397e470395c2b8020c2f1 Mon Sep 17 00:00:00 2001 From: Paul Bottein Date: Fri, 24 Jul 2026 22:11:42 +0200 Subject: [PATCH 2/2] feat(a01): add push listener for unsolicited device state --- roborock/devices/traits/a01/__init__.py | 3 --- tests/devices/traits/a01/test_init.py | 3 +-- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/roborock/devices/traits/a01/__init__.py b/roborock/devices/traits/a01/__init__.py index 739ef30f3..a1cef64e6 100644 --- a/roborock/devices/traits/a01/__init__.py +++ b/roborock/devices/traits/a01/__init__.py @@ -138,9 +138,6 @@ def convert_zeo_value(protocol_value: RoborockZeoProtocol, value: Any) -> Any: return None -# RoborockDyadDataProtocol._missing_ maps any unknown code to its first member -# instead of raising, so incoming data points must be checked against this set -# before being converted to a protocol. _DYAD_PROTOCOL_VALUES = frozenset(protocol.value for protocol in RoborockDyadDataProtocol) diff --git a/tests/devices/traits/a01/test_init.py b/tests/devices/traits/a01/test_init.py index ecdb92474..20130d35f 100644 --- a/tests/devices/traits/a01/test_init.py +++ b/tests/devices/traits/a01/test_init.py @@ -141,7 +141,6 @@ async def test_dyad_add_listener(dyad_api: DyadApi, fake_channel: FakeChannel): received: list[dict[RoborockDyadDataProtocol, Any]] = [] unsub = await dyad_api.add_listener(received.append) - # 999 is not a known protocol: it must be skipped, not mapped to the first enum member. fake_channel.notify_subscribers(build_a01_message({206: 3, 209: 80, 216: 0, 999: 1})) assert received == [ @@ -154,7 +153,7 @@ async def test_dyad_add_listener(dyad_api: DyadApi, fake_channel: FakeChannel): unsub() fake_channel.notify_subscribers(build_a01_message({206: 1})) - assert len(received) == 1 # no callback fires after unsubscribing + assert len(received) == 1 async def test_zeo_api_query_values(zeo_api: ZeoApi, fake_channel: FakeChannel):