File: test_device_manager.py

package info (click to toggle)
python-roborock 4.12.0-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 2,480 kB
  • sloc: python: 16,602; makefile: 17; sh: 6
file content (422 lines) | stat: -rw-r--r-- 16,098 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
"""Tests for the DeviceManager class."""

import asyncio
import datetime
from collections.abc import Generator, Iterator
from typing import Any
from unittest.mock import AsyncMock, Mock, patch

import pytest
import syrupy

from roborock.data import HomeData, UserData
from roborock.devices.cache import InMemoryCache
from roborock.devices.device import RoborockDevice
from roborock.devices.device_manager import UserParams, create_device_manager, create_web_api_wrapper
from roborock.exceptions import RoborockException
from tests import mock_data

USER_DATA = UserData.from_dict(mock_data.USER_DATA)
USER_PARAMS = UserParams(username="test_user", user_data=USER_DATA)
NETWORK_INFO = mock_data.NETWORK_INFO


@pytest.fixture(autouse=True, name="mqtt_session")
def setup_mqtt_session() -> Generator[Mock, None, None]:
    """Fixture to set up the MQTT session for the tests."""
    with patch("roborock.devices.device_manager.create_lazy_mqtt_session") as mock_create_session:
        yield mock_create_session


@pytest.fixture(autouse=True, name="mock_rpc_channel")
def rpc_channel_fixture() -> AsyncMock:
    """Fixture to set up the channel for tests."""
    return AsyncMock()


@pytest.fixture(autouse=True)
async def discover_features_fixture(
    mock_rpc_channel: AsyncMock,
) -> None:
    """Fixture to handle device feature discovery."""
    mock_rpc_channel.send_command.side_effect = [
        [mock_data.APP_GET_INIT_STATUS],
        mock_data.STATUS,
    ]


@pytest.fixture(autouse=True)
def channel_fixture(mock_rpc_channel: AsyncMock) -> Generator[Mock, None, None]:
    """Fixture to set up the local session for the tests."""
    with patch("roborock.devices.device_manager.create_v1_channel") as mock_channel:
        mock_unsub = Mock()
        mock_channel.return_value.subscribe = AsyncMock()
        mock_channel.return_value.subscribe.return_value = mock_unsub
        mock_channel.return_value.rpc_channel = mock_rpc_channel
        yield mock_channel


@pytest.fixture(autouse=True)
def mock_sleep() -> Generator[None, None, None]:
    """Mock sleep logic to speed up tests."""
    sleep_time = datetime.timedelta(seconds=0.001)
    with (
        patch("roborock.devices.device.MIN_BACKOFF_INTERVAL", sleep_time),
        patch("roborock.devices.device.MAX_BACKOFF_INTERVAL", sleep_time),
    ):
        yield


@pytest.fixture(name="channel_exception")
def channel_failure_exception_fixture(mock_rpc_channel: AsyncMock) -> Exception:
    """Fixture that provides the exception to be raised by the failing channel."""
    return RoborockException("Connection failed")


@pytest.fixture(name="channel_failure")
def channel_failure_fixture(mock_rpc_channel: AsyncMock, channel_exception: Exception) -> Generator[Mock, None, None]:
    """Fixture that makes channel subscribe fail."""
    with patch("roborock.devices.device_manager.create_v1_channel") as mock_channel:
        mock_channel.return_value.subscribe = AsyncMock(side_effect=channel_exception)
        mock_channel.return_value.is_connected = False
        mock_channel.return_value.rpc_channel = mock_rpc_channel
        yield mock_channel


@pytest.fixture(name="home_data_no_devices")
def home_data_no_devices_fixture() -> Iterator[HomeData]:
    """Mock home data API that returns no devices."""
    with patch("roborock.devices.device_manager.UserWebApiClient.get_home_data") as mock_home_data:
        home_data = HomeData(
            id=1,
            name="Test Home",
            devices=[],
            products=[],
        )
        mock_home_data.return_value = home_data
        yield home_data


@pytest.fixture(name="home_data")
def home_data_fixture() -> Iterator[HomeData]:
    """Mock home data API that returns devices."""
    with patch("roborock.devices.device_manager.UserWebApiClient.get_home_data") as mock_home_data:
        home_data = HomeData.from_dict(mock_data.HOME_DATA_RAW)
        mock_home_data.return_value = home_data
        yield home_data


async def test_no_devices(home_data_no_devices: HomeData) -> None:
    """Test the DeviceManager created with no devices returned from the API."""

    device_manager = await create_device_manager(USER_PARAMS)
    devices = await device_manager.get_devices()
    assert devices == []


async def test_with_device(home_data: HomeData) -> None:
    """Test the DeviceManager created with devices returned from the API."""
    device_manager = await create_device_manager(USER_PARAMS)
    devices = await device_manager.get_devices()
    assert len(devices) == 1
    assert devices[0].duid == "abc123"
    assert devices[0].name == "Roborock S7 MaxV"

    device = await device_manager.get_device("abc123")
    assert device is not None
    assert device.duid == "abc123"
    assert device.name == "Roborock S7 MaxV"

    await device_manager.close()


async def test_get_non_existent_device(home_data: HomeData) -> None:
    """Test getting a non-existent device."""
    device_manager = await create_device_manager(USER_PARAMS)
    device = await device_manager.get_device("non_existent_duid")
    assert device is None
    await device_manager.close()


async def test_create_home_data_api_exception() -> None:
    """Test that exceptions from the home data API are propagated through the wrapper."""

    with patch("roborock.devices.device_manager.RoborockApiClient.get_home_data_v3") as mock_get_home_data:
        mock_get_home_data.side_effect = RoborockException("Test exception")
        user_params = UserParams(username="test_user", user_data=USER_DATA)
        api = create_web_api_wrapper(user_params)

        with pytest.raises(RoborockException, match="Test exception"):
            await api.get_home_data()


@pytest.mark.parametrize(("prefer_cache", "expected_call_count"), [(True, 1), (False, 2)])
async def test_cache_logic(prefer_cache: bool, expected_call_count: int) -> None:
    """Test that the cache logic works correctly."""
    call_count = 0

    async def mock_home_data_with_counter(*args, **kwargs) -> HomeData:
        nonlocal call_count
        call_count += 1
        return HomeData.from_dict(mock_data.HOME_DATA_RAW)

    # First call happens during create_device_manager initialization
    with patch(
        "roborock.devices.device_manager.RoborockApiClient.get_home_data_v3",
        side_effect=mock_home_data_with_counter,
    ):
        device_manager = await create_device_manager(USER_PARAMS, cache=InMemoryCache())
        assert call_count == 1

        # Second call should use cache, not increment call_count
        devices2 = await device_manager.discover_devices(prefer_cache=prefer_cache)
        assert call_count == expected_call_count
        assert len(devices2) == 1

        await device_manager.close()
        assert len(devices2) == 1

        # Ensure closing again works without error
        await device_manager.close()


async def test_home_data_api_fails_with_cache_fallback() -> None:
    """Test that home data exceptions may still fall back to use the cache when available."""

    cache = InMemoryCache()
    cache_data = await cache.get()
    cache_data.home_data = HomeData.from_dict(mock_data.HOME_DATA_RAW)
    await cache.set(cache_data)

    with patch(
        "roborock.devices.device_manager.RoborockApiClient.get_home_data_v3",
        side_effect=RoborockException("Test exception"),
    ):
        # This call will skip the API and use the cache
        device_manager = await create_device_manager(USER_PARAMS, cache=cache)

        # This call will hit the API since we're not preferring the cache
        # but will fallback to the cache data on exception
        devices2 = await device_manager.discover_devices(prefer_cache=False)
        assert len(devices2) == 1

        await device_manager.close()


async def test_ready_callback(home_data: HomeData) -> None:
    """Test that the ready callback is invoked when a device connects."""
    ready_devices: list[RoborockDevice] = []
    device_manager = await create_device_manager(USER_PARAMS, ready_callback=ready_devices.append)

    # Callback should be called for the discovered device
    assert len(ready_devices) == 1
    device = ready_devices[0]
    assert device.duid == "abc123"

    # Verify that adding a ready callback to an already connected device will
    # invoke the callback immediately.
    more_ready_device: list[RoborockDevice] = []
    device.add_ready_callback(more_ready_device.append)
    assert len(more_ready_device) == 1
    assert more_ready_device[0].duid == "abc123"

    await device_manager.close()


@pytest.mark.parametrize(
    ("channel_exception"),
    [
        RoborockException("Connection failed"),
    ],
)
async def test_start_connect_failure(home_data: HomeData, channel_failure: Mock, mock_sleep: Mock) -> None:
    """Test that start_connect retries when connection fails."""
    ready_devices: list[RoborockDevice] = []
    device_manager = await create_device_manager(USER_PARAMS, ready_callback=ready_devices.append)
    devices = await device_manager.get_devices()

    # The device should attempt to connect in the background at least once
    # by the time this function returns.
    subscribe_mock = channel_failure.return_value.subscribe
    assert subscribe_mock.call_count > 0

    # Device should exist but not be connected
    assert len(devices) == 1
    assert not devices[0].is_connected
    assert not ready_devices

    # Verify retry attempts
    assert channel_failure.return_value.subscribe.call_count >= 1

    # Reset the mock channel so that it succeeds on the next attempt
    mock_unsub = Mock()
    subscribe_mock = AsyncMock()
    subscribe_mock.return_value = mock_unsub
    channel_failure.return_value.subscribe = subscribe_mock
    channel_failure.return_value.is_connected = True

    # Wait for the device to attempt to connect again
    attempts = 0
    while subscribe_mock.call_count < 1:
        await asyncio.sleep(0.01)
        attempts += 1
        assert attempts < 10, "Device did not connect after multiple attempts"

    assert devices[0].is_connected
    assert ready_devices
    assert len(ready_devices) == 1

    await device_manager.close()
    assert mock_unsub.call_count == 1


async def test_rediscover_devices(mock_rpc_channel: AsyncMock) -> None:
    """Test that we can discover devices multiple times and discover new devices."""
    raw_devices: list[dict[str, Any]] = mock_data.HOME_DATA_RAW["devices"]
    assert len(raw_devices) > 0
    raw_device_1 = raw_devices[0]

    home_data_responses = [
        HomeData.from_dict(mock_data.HOME_DATA_RAW),
        # New device added on second call. We make a copy and updated fields to simulate
        # a new device.
        HomeData.from_dict(
            {
                **mock_data.HOME_DATA_RAW,
                "devices": [
                    raw_device_1,
                    {
                        **raw_device_1,
                        "duid": "new_device_duid",
                        "name": "New Device",
                        "model": "roborock.newmodel.v1",
                        "mac": "00:11:22:33:44:55",
                    },
                ],
            }
        ),
    ]

    mock_rpc_channel.send_command.side_effect = [
        [mock_data.APP_GET_INIT_STATUS],
        mock_data.STATUS,
        # Device #2
        [mock_data.APP_GET_INIT_STATUS],
        mock_data.STATUS,
    ]

    async def mock_home_data_with_counter(*args, **kwargs) -> HomeData:
        nonlocal home_data_responses
        return home_data_responses.pop(0)

    # First call happens during create_device_manager initialization
    with patch(
        "roborock.devices.device_manager.RoborockApiClient.get_home_data_v3",
        side_effect=mock_home_data_with_counter,
    ):
        device_manager = await create_device_manager(USER_PARAMS, cache=InMemoryCache())
        assert len(await device_manager.get_devices()) == 1

        # Second call should use cache and does not add new device
        await device_manager.discover_devices(prefer_cache=True)
        assert len(await device_manager.get_devices()) == 1

        # Third call should fetch new home data and add the new device
        await device_manager.discover_devices(prefer_cache=False)
        assert len(await device_manager.get_devices()) == 2

        # Verify the two devices exist with correct data
        device_1 = await device_manager.get_device("abc123")
        assert device_1 is not None
        assert device_1.name == "Roborock S7 MaxV"

        new_device = await device_manager.get_device("new_device_duid")
        assert new_device
        assert new_device.name == "New Device"

        # Ensure closing again works without error
        await device_manager.close()


@pytest.mark.parametrize(
    ("channel_exception"),
    [
        Exception("Unexpected error"),
    ],
)
async def test_start_connect_unexpected_error(home_data: HomeData, channel_failure: Mock, mock_sleep: Mock) -> None:
    """Test that some unexpected errors from start_connect are propagated."""
    with pytest.raises(Exception, match="Unexpected error"):
        await create_device_manager(USER_PARAMS)


async def test_diagnostics_collection(home_data: HomeData, snapshot: syrupy.SnapshotAssertion) -> None:
    """Test that diagnostics are collected correctly in the DeviceManager."""
    device_manager = await create_device_manager(USER_PARAMS)
    devices = await device_manager.get_devices()
    assert len(devices) == 1

    diagnostics = device_manager.diagnostic_data()
    assert diagnostics is not None
    diagnostics_data = diagnostics.get("diagnostics")
    assert diagnostics_data
    assert diagnostics_data.get("discover_devices") == 1
    assert diagnostics_data.get("fetch_home_data") == 1

    assert snapshot == diagnostics

    await device_manager.close()


async def test_unsupported_protocol_version() -> None:
    """Test the DeviceManager with some supported and unsupported product IDs."""
    with patch("roborock.devices.device_manager.UserWebApiClient.get_home_data") as mock_home_data:
        home_data = HomeData.from_dict(
            {
                "id": 1,
                "name": "Test Home",
                "devices": [
                    {
                        "duid": "device-uid-1",
                        "name": "Device 1",
                        "pv": "1.0",
                        "productId": "product-id-1",
                        "localKey": mock_data.LOCAL_KEY,
                    },
                    {
                        "duid": "device-uid-2",
                        "name": "Device 2",
                        "pv": "unknown-pv",  # Fake new protocol version we've never seen
                        "productId": "product-id-2",
                        "localKey": mock_data.LOCAL_KEY,
                    },
                ],
                "products": [
                    {
                        "id": "product-id-1",
                        "name": "Roborock S7 MaxV",
                        "model": "roborock.vacuum.a27",
                        "category": "robot.vacuum.cleaner",
                    },
                    {
                        "id": "product-id-2",
                        "name": "New Roborock Model",
                        "model": "roborock.vacuum.newmodel",
                        "category": "robot.vacuum.cleaner",
                    },
                ],
            }
        )
        mock_home_data.return_value = home_data

        device_manager = await create_device_manager(USER_PARAMS)
        # Only the supported device should be created. The other device is ignored
        devices = await device_manager.get_devices()
        assert [device.duid for device in devices] == ["device-uid-1"]

        # Verify diagnostics
        diagnostics = device_manager.diagnostic_data()
        diagnostics_data = diagnostics.get("diagnostics")
        assert diagnostics_data
        assert diagnostics_data.get("supported_devices") == {"1.0": 1}
        assert diagnostics_data.get("unsupported_devices") == {"unknown-pv": 1}