Skip to content

Commit 3994110

Browse files
allenporterCopilot
andauthored
feat: Connect to devices asynchronously (#588)
* feat: Connect to devices asynchronously * chore: Update roborock/devices/device.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * fix: Update roborock/devices/device.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * chore: Update pydoc and formatting * chore: update comments to clarify close call * chore: update typing * chore: Fix lint errors --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent 3296a3a commit 3994110

File tree

3 files changed

+111
-1
lines changed

3 files changed

+111
-1
lines changed

roborock/devices/device.py

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,15 @@
44
until the API is stable.
55
"""
66

7+
import asyncio
8+
import datetime
79
import logging
810
from abc import ABC
911
from collections.abc import Callable, Mapping
1012
from typing import Any, TypeVar, cast
1113

1214
from roborock.data import HomeDataDevice, HomeDataProduct
15+
from roborock.exceptions import RoborockException
1316
from roborock.roborock_message import RoborockMessage
1417

1518
from .channel import Channel
@@ -22,6 +25,11 @@
2225
"RoborockDevice",
2326
]
2427

28+
# Exponential backoff parameters
29+
MIN_BACKOFF_INTERVAL = datetime.timedelta(seconds=10)
30+
MAX_BACKOFF_INTERVAL = datetime.timedelta(minutes=30)
31+
BACKOFF_MULTIPLIER = 1.5
32+
2533

2634
class RoborockDevice(ABC, TraitsMixin):
2735
"""A generic channel for establishing a connection with a Roborock device.
@@ -54,6 +62,7 @@ def __init__(
5462
self._device_info = device_info
5563
self._product = product
5664
self._channel = channel
65+
self._connect_task: asyncio.Task[None] | None = None
5766
self._unsub: Callable[[], None] | None = None
5867

5968
@property
@@ -98,6 +107,38 @@ def is_local_connected(self) -> bool:
98107
"""
99108
return self._channel.is_local_connected
100109

110+
def start_connect(self) -> None:
111+
"""Start a background task to connect to the device.
112+
113+
This will attempt to connect to the device using the appropriate protocol
114+
channel. If the connection fails, it will retry with exponential backoff.
115+
116+
Once connected, the device will remain connected until `close()` is
117+
called. The device will automatically attempt to reconnect if the connection
118+
is lost.
119+
"""
120+
121+
async def connect_loop() -> None:
122+
backoff = MIN_BACKOFF_INTERVAL
123+
try:
124+
while True:
125+
try:
126+
await self.connect()
127+
return
128+
except RoborockException as e:
129+
_LOGGER.info("Failed to connect to device %s: %s", self.name, e)
130+
_LOGGER.info(
131+
"Retrying connection to device %s in %s seconds", self.name, backoff.total_seconds()
132+
)
133+
await asyncio.sleep(backoff.total_seconds())
134+
backoff = min(backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_INTERVAL)
135+
except asyncio.CancelledError:
136+
_LOGGER.info("connect_loop for device %s was cancelled", self.name)
137+
# Clean exit on cancellation
138+
return
139+
140+
self._connect_task = asyncio.create_task(connect_loop())
141+
101142
async def connect(self) -> None:
102143
"""Connect to the device using the appropriate protocol channel."""
103144
if self._unsub:
@@ -107,6 +148,12 @@ async def connect(self) -> None:
107148

108149
async def close(self) -> None:
109150
"""Close all connections to the device."""
151+
if self._connect_task:
152+
self._connect_task.cancel()
153+
try:
154+
await self._connect_task
155+
except asyncio.CancelledError:
156+
pass
110157
if self._unsub:
111158
self._unsub()
112159
self._unsub = None

roborock/devices/device_manager.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ async def discover_devices(self) -> list[RoborockDevice]:
8686
if duid in self._devices:
8787
continue
8888
new_device = self._device_creator(home_data, device, product)
89-
await new_device.connect()
89+
new_device.start_connect()
9090
new_devices[duid] = new_device
9191

9292
self._devices.update(new_devices)

tests/devices/test_device_manager.py

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
"""Tests for the DeviceManager class."""
22

3+
import asyncio
4+
import datetime
35
from collections.abc import Generator, Iterator
46
from unittest.mock import AsyncMock, Mock, patch
57

@@ -34,6 +36,26 @@ def channel_fixture() -> Generator[Mock, None, None]:
3436
yield mock_channel
3537

3638

39+
@pytest.fixture(autouse=True)
40+
def mock_sleep() -> Generator[None, None, None]:
41+
"""Mock sleep logic to speed up tests."""
42+
sleep_time = datetime.timedelta(seconds=0.001)
43+
with (
44+
patch("roborock.devices.device.MIN_BACKOFF_INTERVAL", sleep_time),
45+
patch("roborock.devices.device.MAX_BACKOFF_INTERVAL", sleep_time),
46+
):
47+
yield
48+
49+
50+
@pytest.fixture(name="channel_failure")
51+
def channel_failure_fixture() -> Generator[Mock, None, None]:
52+
"""Fixture that makes channel subscribe fail."""
53+
with patch("roborock.devices.device_manager.create_v1_channel") as mock_channel:
54+
mock_channel.return_value.subscribe = AsyncMock(side_effect=RoborockException("Connection failed"))
55+
mock_channel.return_value.is_connected = False
56+
yield mock_channel
57+
58+
3759
@pytest.fixture(name="home_data_no_devices")
3860
def home_data_no_devices_fixture() -> Iterator[HomeData]:
3961
"""Mock home data API that returns no devices."""
@@ -126,4 +148,45 @@ async def mock_home_data_with_counter(*args, **kwargs) -> HomeData:
126148
await device_manager.close()
127149
assert len(devices2) == 1
128150

151+
# Ensure closing again works without error
129152
await device_manager.close()
153+
154+
155+
async def test_start_connect_failure(home_data: HomeData, channel_failure: Mock, mock_sleep: Mock) -> None:
156+
"""Test that start_connect retries when connection fails."""
157+
device_manager = await create_device_manager(USER_PARAMS)
158+
devices = await device_manager.get_devices()
159+
160+
# Wait for the device to attempt to connect
161+
attempts = 0
162+
subscribe_mock = channel_failure.return_value.subscribe
163+
while subscribe_mock.call_count < 1:
164+
await asyncio.sleep(0.01)
165+
attempts += 1
166+
assert attempts < 10, "Device did not connect after multiple attempts"
167+
168+
# Device should exist but not be connected
169+
assert len(devices) == 1
170+
assert not devices[0].is_connected
171+
172+
# Verify retry attempts
173+
assert channel_failure.return_value.subscribe.call_count >= 1
174+
175+
# Reset the mock channel so that it succeeds on the next attempt
176+
mock_unsub = Mock()
177+
subscribe_mock = AsyncMock()
178+
subscribe_mock.return_value = mock_unsub
179+
channel_failure.return_value.subscribe = subscribe_mock
180+
channel_failure.return_value.is_connected = True
181+
182+
# Wait for the device to attempt to connect again
183+
attempts = 0
184+
while subscribe_mock.call_count < 1:
185+
await asyncio.sleep(0.01)
186+
attempts += 1
187+
assert attempts < 10, "Device did not connect after multiple attempts"
188+
189+
assert devices[0].is_connected
190+
191+
await device_manager.close()
192+
assert mock_unsub.call_count == 1

0 commit comments

Comments
 (0)