File: test_local_channel.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 (469 lines) | stat: -rw-r--r-- 17,637 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
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
"""Tests for the LocalChannel class."""

import asyncio
import json
from collections.abc import Generator
from unittest.mock import AsyncMock, Mock, patch

import pytest

from roborock.devices.transport.local_channel import LocalChannel, LocalChannelParams
from roborock.exceptions import RoborockConnectionException, RoborockException
from roborock.protocol import create_local_decoder, create_local_encoder
from roborock.protocols.v1_protocol import LocalProtocolVersion
from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol

TEST_HOST = "192.168.1.100"
TEST_LOCAL_KEY = "local_key"
TEST_PORT = 58867

TEST_REQUEST = RoborockMessage(
    protocol=RoborockMessageProtocol.RPC_REQUEST,
    payload=json.dumps({"dps": {"101": json.dumps({"id": 12345, "method": "get_status"})}}).encode(),
)
TEST_RESPONSE = RoborockMessage(
    protocol=RoborockMessageProtocol.RPC_RESPONSE,
    payload=json.dumps({"dps": {"102": json.dumps({"id": 12345, "result": {"state": "cleaning"}})}}).encode(),
)
TEST_REQUEST2 = RoborockMessage(
    protocol=RoborockMessageProtocol.RPC_REQUEST,
    payload=json.dumps({"dps": {"101": json.dumps({"id": 54321, "method": "get_status"})}}).encode(),
)
TEST_RESPONSE2 = RoborockMessage(
    protocol=RoborockMessageProtocol.RPC_RESPONSE,
    payload=json.dumps({"dps": {"102": json.dumps({"id": 54321, "result": {"state": "cleaning"}})}}).encode(),
)
ENCODER = create_local_encoder(TEST_LOCAL_KEY)
DECODER = create_local_decoder(TEST_LOCAL_KEY)


@pytest.fixture(name="mock_transport")
def setup_mock_transport() -> Mock:
    """Mock transport for testing."""
    transport = Mock()
    transport.write = Mock()
    transport.close = Mock()
    return transport


@pytest.fixture(name="mock_loop")
def setup_mock_loop(mock_transport: Mock) -> Generator[Mock, None, None]:
    """Mock event loop for testing."""
    loop = Mock()
    loop.create_connection = AsyncMock(return_value=(mock_transport, Mock()))

    with patch("roborock.devices.transport.local_channel.get_running_loop", return_value=loop):
        yield loop


def create_test_local_channel() -> LocalChannel:
    """Helper to create a LocalChannel for testing."""
    return LocalChannel(host=TEST_HOST, local_key=TEST_LOCAL_KEY, device_uid="test_duid")


@pytest.fixture(name="local_channel")
async def setup_local_channel_with_hello_mock() -> LocalChannel:
    """Fixture to set up the local channel with automatic hello mocking."""
    channel = create_test_local_channel()

    async def mock_do_hello(_: LocalProtocolVersion):
        """Mock _do_hello to return successful params without sending actual request."""
        return LocalChannelParams(
            local_key=channel._params.local_key, connect_nonce=channel._params.connect_nonce, ack_nonce=54321
        )

    # Replace the _do_hello method
    setattr(channel, "_do_hello", mock_do_hello)

    return channel


@pytest.fixture(name="received_messages")
async def setup_subscribe_callback(local_channel: LocalChannel) -> list[RoborockMessage]:
    """Fixture to record messages received by the subscriber."""
    messages: list[RoborockMessage] = []
    await local_channel.subscribe(messages.append)
    return messages


async def test_successful_connection(local_channel: LocalChannel, mock_loop: Mock, mock_transport: Mock) -> None:
    """Test successful connection to device."""
    await local_channel.connect()

    mock_loop.create_connection.assert_called_once()
    call_args = mock_loop.create_connection.call_args
    assert call_args[0][1] == TEST_HOST
    assert call_args[0][2] == TEST_PORT
    assert local_channel._is_connected is True


async def test_connection_failure(local_channel: LocalChannel, mock_loop: Mock) -> None:
    """Test connection failure handling."""
    mock_loop.create_connection.side_effect = OSError("Connection failed")

    with pytest.raises(RoborockConnectionException, match="Failed to connect to 192.168.1.100:58867"):
        await local_channel.connect()

    assert local_channel._is_connected is False


async def test_close_connection(local_channel: LocalChannel, mock_loop: Mock, mock_transport: Mock) -> None:
    """Test closing the connection."""
    await local_channel.connect()
    local_channel.close()

    mock_transport.close.assert_called_once()
    assert local_channel._is_connected is False


async def test_close_without_connection(local_channel: LocalChannel) -> None:
    """Test closing when not connected."""
    local_channel.close()
    assert local_channel._is_connected is False


async def test_publish_not_connected(local_channel: LocalChannel) -> None:
    """Test sending command when not connected raises exception."""
    with pytest.raises(RoborockConnectionException, match="Not connected to device"):
        await local_channel.publish(TEST_REQUEST)


async def test_successful_command_response(local_channel: LocalChannel, mock_loop: Mock, mock_transport: Mock) -> None:
    """Test successful command sending and response handling."""
    await local_channel.connect()

    # Send command in background task
    await local_channel.publish(TEST_REQUEST)
    await asyncio.sleep(0.01)  # yield

    # Simulate receiving response via the protocol callback
    local_channel._data_received(ENCODER(TEST_RESPONSE))
    await asyncio.sleep(0.01)  # yield

    # Verify command was sent
    mock_transport.write.assert_called_once()
    sent_data = mock_transport.write.call_args[0][0]
    decoded_sent = next(iter(DECODER(sent_data)))
    assert decoded_sent == TEST_REQUEST


async def test_message_decode_error(
    local_channel: LocalChannel, caplog: pytest.LogCaptureFixture, received_messages: list[RoborockMessage]
) -> None:
    """Test handling of message decode errors."""
    local_channel._data_received(b"invalid_payload")
    await asyncio.sleep(0.01)  # yield

    assert received_messages == []


async def test_subscribe_callback(
    local_channel: LocalChannel, received_messages: list[RoborockMessage], mock_loop: Mock
) -> None:
    """Test that subscribe callback receives all messages."""
    await local_channel.connect()

    # Send some messages without an RPC
    local_channel._data_received(ENCODER(TEST_RESPONSE))
    local_channel._data_received(ENCODER(TEST_RESPONSE2))
    await asyncio.sleep(0.01)  # yield

    assert received_messages == [TEST_RESPONSE, TEST_RESPONSE2]


async def test_subscribe_callback_exception_handling(
    local_channel: LocalChannel, mock_loop: Mock, caplog: pytest.LogCaptureFixture
) -> None:
    """Test that exceptions in subscriber callbacks are handled gracefully."""

    def failing_callback(message: RoborockMessage) -> None:
        raise ValueError("Test exception")

    await local_channel.subscribe(failing_callback)
    await local_channel.connect()

    # Send message that will cause callback to fail
    local_channel._data_received(ENCODER(TEST_RESPONSE))
    await asyncio.sleep(0.01)  # yield

    # Should log the exception but not crash
    assert any("Uncaught error in callback 'failing_callback'" in record.message for record in caplog.records)


async def test_unsubscribe(local_channel: LocalChannel, mock_loop: Mock) -> None:
    """Test unsubscribing from messages."""
    messages: list[RoborockMessage] = []
    unsubscribe = await local_channel.subscribe(messages.append)
    await local_channel.connect()

    # Send message while subscribed
    local_channel._data_received(ENCODER(TEST_RESPONSE))
    await asyncio.sleep(0.01)  # yield
    assert len(messages) == 1

    # Unsubscribe and send another message
    unsubscribe()
    local_channel._data_received(ENCODER(TEST_RESPONSE2))
    await asyncio.sleep(0.01)  # yield

    # Should still have only one message
    assert len(messages) == 1


async def test_connection_lost_callback(
    local_channel: LocalChannel, mock_loop: Mock, caplog: pytest.LogCaptureFixture
) -> None:
    """Test connection lost callback handling."""
    await local_channel.connect()

    # Simulate connection loss
    test_exception = OSError("Connection lost")
    local_channel._connection_lost(test_exception)

    assert local_channel._is_connected is False
    assert local_channel._transport is None


async def test_connection_lost_without_exception(
    local_channel: LocalChannel, mock_loop: Mock, caplog: pytest.LogCaptureFixture
) -> None:
    """Test connection lost callback without exception."""
    await local_channel.connect()

    # Simulate connection loss without exception
    local_channel._connection_lost(None)

    assert local_channel._is_connected is False
    assert local_channel._transport is None


async def test_hello_fallback_to_l01_protocol(mock_loop: Mock, mock_transport: Mock) -> None:
    """Test that when first hello() message fails (V1) but second succeeds (L01), we use L01."""

    # Create a channel without the automatic hello mocking
    channel = create_test_local_channel()

    # Mock _do_hello to fail for V1 but succeed for L01
    async def mock_do_hello(local_protocol_version: LocalProtocolVersion) -> LocalChannelParams | None:
        if local_protocol_version == LocalProtocolVersion.V1:
            # First attempt (V1) fails - return None to simulate failure
            return None
        elif local_protocol_version == LocalProtocolVersion.L01:
            # Second attempt (L01) succeeds
            return LocalChannelParams(
                local_key=channel._params.local_key, connect_nonce=channel._params.connect_nonce, ack_nonce=54321
            )
        return None

    # Replace the _do_hello method
    setattr(channel, "_do_hello", mock_do_hello)

    # Connect and verify L01 protocol is used
    await channel.connect()

    # Verify that the channel is using L01 protocol
    assert channel._local_protocol_version == LocalProtocolVersion.L01
    assert channel._params is not None
    assert channel._params.ack_nonce == 54321
    assert channel._is_connected is True


async def test_hello_success_with_v1_protocol_first(mock_loop: Mock, mock_transport: Mock) -> None:
    """Test that when V1 protocol succeeds on first attempt, we use V1."""

    # Create a channel without the automatic hello mocking
    channel = create_test_local_channel()
    # Clear cached protocol to ensure V1 is tried first
    channel._local_protocol_version = None

    # Mock _do_hello to succeed for V1 on first attempt
    async def mock_do_hello(local_protocol_version: LocalProtocolVersion) -> LocalChannelParams | None:
        if local_protocol_version == LocalProtocolVersion.V1:
            # V1 succeeds on first attempt
            return LocalChannelParams(
                local_key=channel._params.local_key, connect_nonce=channel._params.connect_nonce, ack_nonce=67890
            )
        elif local_protocol_version == LocalProtocolVersion.L01:
            # L01 would succeed but we shouldn't reach it
            return LocalChannelParams(
                local_key=channel._params.local_key, connect_nonce=channel._params.connect_nonce, ack_nonce=99999
            )
        return None

    # Replace the _do_hello method
    setattr(channel, "_do_hello", mock_do_hello)

    # Connect and verify V1 protocol is used
    await channel.connect()

    # Verify that the channel is using V1 protocol
    assert channel._local_protocol_version == LocalProtocolVersion.V1
    assert channel._params is not None
    assert channel._params.ack_nonce == 67890
    assert channel._is_connected is True


async def test_hello_both_protocols_fail(mock_loop: Mock, mock_transport: Mock) -> None:
    """Test that when both V1 and L01 protocols fail, connection fails."""

    # Create a channel without the automatic hello mocking
    channel = create_test_local_channel()

    # Mock _do_hello to fail for both protocols
    async def mock_do_hello(_: LocalProtocolVersion) -> LocalChannelParams | None:
        # Both protocols fail
        return None

    # Replace the _do_hello method
    setattr(channel, "_do_hello", mock_do_hello)

    # Connect should raise an exception
    with pytest.raises(RoborockException, match="Failed to connect to device with any known protocol"):
        await channel.connect()

    # Verify that the channel is not connected and cleaned up
    assert channel._is_connected is False
    assert channel._transport is None


async def test_hello_preferred_protocol_version_ordering(mock_loop: Mock, mock_transport: Mock) -> None:
    """Test that preferred protocol version is tried first."""

    # Create a channel with preferred L01 protocol
    channel = create_test_local_channel()
    channel._local_protocol_version = LocalProtocolVersion.L01

    # Track which protocols were attempted and in what order
    attempted_protocols: list[LocalProtocolVersion] = []

    # Mock _do_hello to track attempts and succeed on L01
    async def mock_do_hello(local_protocol_version: LocalProtocolVersion) -> LocalChannelParams | None:
        attempted_protocols.append(local_protocol_version)
        if local_protocol_version == LocalProtocolVersion.L01:
            # L01 succeeds
            return LocalChannelParams(
                local_key=channel._params.local_key, connect_nonce=channel._params.connect_nonce, ack_nonce=11111
            )
        return None

    # Replace the _do_hello method
    setattr(channel, "_do_hello", mock_do_hello)

    # Connect and verify L01 is tried first
    await channel.connect()

    # Verify that L01 was tried first (preferred version)
    assert attempted_protocols == [LocalProtocolVersion.L01]
    assert channel._local_protocol_version == LocalProtocolVersion.L01
    assert channel._params is not None
    assert channel._params.ack_nonce == 11111
    assert channel._is_connected is True


async def test_keep_alive_task_created_on_connect(local_channel: LocalChannel, mock_loop: Mock) -> None:
    """Test that _keep_alive_task is created when connect() is called."""
    # Before connecting, task should be None
    assert local_channel._keep_alive_task is None

    await local_channel.connect()

    # After connecting, task should be created and not done
    assert local_channel._keep_alive_task is not None
    assert isinstance(local_channel._keep_alive_task, asyncio.Task)
    assert not local_channel._keep_alive_task.done()


async def test_keep_alive_task_canceled_on_close(local_channel: LocalChannel, mock_loop: Mock) -> None:
    """Test that the keep-alive task is properly canceled when close() is called."""
    await local_channel.connect()

    # Verify task exists
    task = local_channel._keep_alive_task
    assert task is not None
    assert not task.done()

    # Close the connection
    local_channel.close()

    # Give the task a moment to be cancelled
    await asyncio.sleep(0.01)

    # Task should be canceled and reset to None
    assert task.cancelled() or task.done()
    assert local_channel._keep_alive_task is None


async def test_keep_alive_task_canceled_on_connection_lost(local_channel: LocalChannel, mock_loop: Mock) -> None:
    """Test that the keep-alive task is properly canceled when _connection_lost() is called."""
    await local_channel.connect()

    # Verify task exists
    task = local_channel._keep_alive_task
    assert task is not None
    assert not task.done()

    # Simulate connection loss
    local_channel._connection_lost(None)

    # Give the task a moment to be cancelled
    await asyncio.sleep(0.01)

    # Task should be canceled and reset to None
    assert task.cancelled() or task.done()
    assert local_channel._keep_alive_task is None


async def test_keep_alive_ping_loop_executes_periodically(local_channel: LocalChannel, mock_loop: Mock) -> None:
    """Test that the ping loop continues to execute periodically while connected."""
    await local_channel.connect()

    # Verify the task is running and connected
    assert local_channel._keep_alive_task is not None
    assert not local_channel._keep_alive_task.done()
    assert local_channel._is_connected


async def test_keep_alive_ping_exceptions_handled_gracefully(
    local_channel: LocalChannel, mock_loop: Mock, caplog: pytest.LogCaptureFixture
) -> None:
    """Test that exceptions in the ping loop are handled gracefully without stopping the loop."""
    from roborock.devices.transport.local_channel import _PING_INTERVAL

    # Set log level to capture DEBUG messages
    caplog.set_level("DEBUG")

    ping_call_count = 0

    # Mock the _ping method to always fail
    async def mock_ping() -> None:
        nonlocal ping_call_count
        ping_call_count += 1
        raise Exception("Test ping failure")

    # Also need to mock asyncio.sleep to avoid waiting the full interval
    original_sleep = asyncio.sleep

    async def mock_sleep(delay: float) -> None:
        # Only sleep briefly for test speed when waiting for ping interval
        if delay >= _PING_INTERVAL:
            await original_sleep(0.01)
        else:
            await original_sleep(delay)

    with patch("asyncio.sleep", side_effect=mock_sleep):
        setattr(local_channel, "_ping", mock_ping)

        await local_channel.connect()

        # Wait for multiple ping attempts
        await original_sleep(0.1)

        # Verify the task is still running despite the exception
        assert local_channel._keep_alive_task is not None
        assert not local_channel._keep_alive_task.done()

        # Verify ping was called at least once
        assert ping_call_count >= 1

        # Verify the exception was logged but didn't crash the loop
        assert any("Keep-alive ping failed" in record.message for record in caplog.records)