From 8006f5e59e0d13143ffdc34f17ddc735059cd146 Mon Sep 17 00:00:00 2001 From: borisalekseev Date: Tue, 8 Sep 2026 01:23:12 +0300 Subject: [PATCH] chore(tests): Cover mqtt e2e scenarios for MqttSession --- .github/workflows/ci.yml | 10 ++ CONTRIBUTING.md | 8 +- compose.yaml | 10 ++ pyproject.toml | 2 + tests/e2e/test_mqtt_broker.py | 240 ++++++++++++++++++++++++++++++++++ uv.lock | 11 ++ 6 files changed, 280 insertions(+), 1 deletion(-) create mode 100644 compose.yaml create mode 100644 tests/e2e/test_mqtt_broker.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d1915057..b560ab46 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -38,6 +38,16 @@ jobs: - "3.11" - "3.14" runs-on: ubuntu-latest + services: + emqx: + image: emqx/emqx:6.2.3 + ports: + - 1888:1883 + options: >- + --health-cmd "/opt/emqx/bin/emqx ctl status" + --health-interval 5s + --health-timeout 25s + --health-retries 5 steps: - uses: actions/checkout@v6 - name: Set up uv diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index bfa8e8b2..b8c0e379 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -53,11 +53,17 @@ pre-commit run --all-files We use `pytest` for testing. Please ensure all tests pass and add new tests for your changes. +MQTT tests require a real EMQX broker at `127.0.0.1:1888`. Start it locally with +Docker Compose before running the tests. CI provides the broker as a GitHub Actions service. + ```bash -# Run tests +docker compose up -d --wait pytest +docker compose down ``` +To run only the real broker tests, use `pytest -m mqtt_broker`. + ## Pull Requests 1. **Create a branch** for your changes. diff --git a/compose.yaml b/compose.yaml new file mode 100644 index 00000000..1e186fa5 --- /dev/null +++ b/compose.yaml @@ -0,0 +1,10 @@ +services: + emqx: + image: emqx/emqx:6.2.3 + ports: + - "127.0.0.1:1888:1883" + healthcheck: + test: ["CMD", "/opt/emqx/bin/emqx", "ctl", "status"] + interval: 5s + timeout: 25s + retries: 5 diff --git a/pyproject.toml b/pyproject.toml index 811d28e0..0f37aeb9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,7 @@ dev = [ "syrupy>=4.9.1,<6", "pdoc>=15.0.4,<17", "pytest-cov>=7.0.0", + "zmqtt>=0.2.0,<0.3.0", ] [tool.hatch.build.targets.sdist] @@ -120,6 +121,7 @@ module = ["roborock.map.proto.*"] ignore_errors = true [tool.pytest.ini_options] +markers = ["mqtt_broker: tests requiring a real MQTT broker"] asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" timeout = 30 diff --git a/tests/e2e/test_mqtt_broker.py b/tests/e2e/test_mqtt_broker.py new file mode 100644 index 00000000..54f77274 --- /dev/null +++ b/tests/e2e/test_mqtt_broker.py @@ -0,0 +1,240 @@ +"""Test the MQTT session and channel over TCP with a real broker.""" + +import asyncio +import json +from collections.abc import AsyncGenerator +from uuid import uuid4 + +import pytest +from zmqtt import MQTTClientV5, QoS, ReconnectConfig, create_client + +from roborock.data import UserData +from roborock.devices.transport.mqtt_channel import MqttChannel +from roborock.mqtt.roborock_session import create_lazy_mqtt_session, create_mqtt_session +from roborock.mqtt.session import MqttParams, MqttQos, MqttSession, MqttSessionException +from roborock.protocol import MessageParser +from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol +from tests.mock_data import LOCAL_KEY, USER_DATA + +pytestmark = pytest.mark.mqtt_broker + + +@pytest.fixture(name="mqtt_params") +def mqtt_params_fixture() -> MqttParams: + """Use the broker exposed by Docker Compose or GitHub Actions.""" + return MqttParams( + host="127.0.0.1", + port=1888, + tls=False, + username="username", + password="password", + timeout=5.0, + ) + + +@pytest.fixture(name="session") +async def session_fixture(mqtt_params: MqttParams) -> AsyncGenerator[MqttSession, None]: + """Create and close the production MQTT session.""" + session = await create_mqtt_session(mqtt_params) + try: + assert session.connected + yield session + finally: + await session.close() + + +@pytest.fixture(name="peer") +async def peer_fixture(mqtt_params: MqttParams) -> AsyncGenerator[MQTTClientV5, None]: + """Represent a device using an independent MQTT client.""" + async with create_client( + mqtt_params.host, + mqtt_params.port, + version="5.0", + mqtt_connect_timeout=mqtt_params.timeout, + reconnect=ReconnectConfig(enabled=False), + ) as peer: + yield peer + + +@pytest.fixture(name="topic") +def topic_fixture() -> str: + """Isolate each test's traffic on the shared broker.""" + return f"roborock-tests/{uuid4().hex}" + + +async def test_receive_message(session: MqttSession, peer: MQTTClientV5, topic: str) -> None: + """Deliver an encrypted Roborock response to the session callback.""" + messages: asyncio.Queue[bytes] = asyncio.Queue() + await session.subscribe(topic, messages.put_nowait) + response = RoborockMessage( + protocol=RoborockMessageProtocol.RPC_RESPONSE, + payload=b'{"result":"ok"}', + seq=123, + ) + payload = MessageParser.build(response, local_key=LOCAL_KEY, prefixed=False) + + await peer.publish(topic, payload) + received = await asyncio.wait_for(messages.get(), timeout=5) + + assert received == payload + parsed, remaining = MessageParser.parse(received, local_key=LOCAL_KEY) + assert not remaining + assert len(parsed) == 1 + assert parsed[0].protocol == response.protocol + assert parsed[0].seq == response.seq + assert parsed[0].payload == response.payload + + +@pytest.mark.parametrize("qos", list(MqttQos)) +async def test_publish_message(session: MqttSession, peer: MQTTClientV5, topic: str, qos: MqttQos) -> None: + """Deliver the full payload and requested QoS to another MQTT client.""" + payload = MessageParser.build( + RoborockMessage(protocol=RoborockMessageProtocol.RPC_REQUEST, payload=b'{"method":"get_status"}'), + local_key=LOCAL_KEY, + prefixed=False, + ) + async with peer.subscribe(topic, qos=QoS.EXACTLY_ONCE) as subscription: + await session.publish(topic, payload, qos=qos) + received = await asyncio.wait_for(subscription.get_message(), timeout=5) + + assert received.topic == topic + assert received.payload == payload + assert received.qos == qos + + +async def test_subscriber_lifecycle(session: MqttSession, peer: MQTTClientV5, topic: str) -> None: + """Fan out messages, remove callbacks, and reuse an idle subscription.""" + first: asyncio.Queue[bytes] = asyncio.Queue() + second: asyncio.Queue[bytes] = asyncio.Queue() + unsub_first = await session.subscribe(topic, first.put_nowait) + unsub_second = await session.subscribe(topic, second.put_nowait) + + await peer.publish(topic, b"both") + assert await asyncio.wait_for(first.get(), timeout=5) == b"both" + assert await asyncio.wait_for(second.get(), timeout=5) == b"both" + + unsub_first() + await peer.publish(topic, b"second only") + assert await asyncio.wait_for(second.get(), timeout=5) == b"second only" + assert first.empty() + + unsub_second() + await session.subscribe(topic, first.put_nowait) + await peer.publish(topic, b"first again") + assert await asyncio.wait_for(first.get(), timeout=5) == b"first again" + assert first.empty() + assert second.empty() + + +async def test_topic_routing(session: MqttSession, peer: MQTTClientV5, topic: str) -> None: + """Keep messages for different devices separate in a shared session.""" + first: asyncio.Queue[bytes] = asyncio.Queue() + second: asyncio.Queue[bytes] = asyncio.Queue() + await session.subscribe(f"{topic}/first", first.put_nowait) + await session.subscribe(f"{topic}/second", second.put_nowait) + + await peer.publish(f"{topic}/first", b"first") + await peer.publish(f"{topic}/second", b"second") + + assert await asyncio.wait_for(first.get(), timeout=5) == b"first" + assert await asyncio.wait_for(second.get(), timeout=5) == b"second" + assert first.empty() + assert second.empty() + + +async def test_restart_restores_subscriptions(session: MqttSession, peer: MQTTClientV5, topic: str) -> None: + """Restore message delivery after the session reconnects.""" + messages: asyncio.Queue[bytes] = asyncio.Queue() + await session.subscribe(topic, messages.put_nowait) + await peer.publish(topic, b"before restart") + assert await asyncio.wait_for(messages.get(), timeout=5) == b"before restart" + + await session.restart() + async with asyncio.timeout(20): + while session.connected: + await asyncio.sleep(0.01) + while not session.connected: + await asyncio.sleep(0.01) + + await peer.publish(topic, b"after restart") + assert await asyncio.wait_for(messages.get(), timeout=5) == b"after restart" + async with peer.subscribe(topic) as subscription: + await session.publish(topic, b"outbound after restart") + received = await asyncio.wait_for(subscription.get_message(), timeout=5) + assert received.payload == b"outbound after restart" + + +async def test_close(session: MqttSession, topic: str) -> None: + """Close a connected session and reject further publications.""" + await session.close() + + assert not session.connected + with pytest.raises(MqttSessionException): + await session.publish(topic, b"closed") + + +async def test_lazy_session_subscribe(mqtt_params: MqttParams, peer: MQTTClientV5, topic: str) -> None: + """Connect a lazy session on its first subscription and receive a message.""" + session = await create_lazy_mqtt_session(mqtt_params) + assert not session.connected + + messages: asyncio.Queue[bytes] = asyncio.Queue() + await session.subscribe(topic, messages.put_nowait) + assert session.connected + + await peer.publish(topic, b"inbound") + assert await asyncio.wait_for(messages.get(), timeout=5) == b"inbound" + + await session.close() + assert not session.connected + + +async def test_lazy_session_publish(mqtt_params: MqttParams, peer: MQTTClientV5, topic: str) -> None: + """Connect a lazy session on its first publication and deliver a message.""" + session = await create_lazy_mqtt_session(mqtt_params) + assert not session.connected + + async with peer.subscribe(topic) as subscription: + await session.publish(topic, b"outbound") + assert session.connected + received = await asyncio.wait_for(subscription.get_message(), timeout=5) + assert received.payload == b"outbound" + + await session.close() + assert not session.connected + + +async def test_channel_request_response(session: MqttSession, mqtt_params: MqttParams, peer: MQTTClientV5) -> None: + """Exchange encrypted commands and responses on the device's Roborock topics.""" + user_data = UserData.from_dict(USER_DATA) + duid = uuid4().hex + channel = MqttChannel(session, duid, LOCAL_KEY, user_data.rriot, mqtt_params) + request_topic = f"rr/m/i/{user_data.rriot.u}/{mqtt_params.username}/{duid}" + response_topic = f"rr/m/o/{user_data.rriot.u}/{mqtt_params.username}/{duid}" + messages: asyncio.Queue[RoborockMessage] = asyncio.Queue() + unsub = await channel.subscribe(messages.put_nowait) + async with peer.subscribe(request_topic) as subscription: + command = RoborockMessage( + protocol=RoborockMessageProtocol.RPC_REQUEST, + payload=json.dumps({"dps": {"101": json.dumps({"id": 123, "method": "get_status"})}}).encode(), + seq=456, + ) + await channel.publish(command) + request = await asyncio.wait_for(subscription.get_message(), timeout=5) + parsed, remaining = MessageParser.parse(request.payload, local_key=LOCAL_KEY) + assert request.topic == request_topic + assert not remaining + assert len(parsed) == 1 + assert parsed[0].protocol == RoborockMessageProtocol.RPC_REQUEST + assert parsed[0].seq == 456 + assert parsed[0].payload == command.payload + + response = RoborockMessage( + protocol=RoborockMessageProtocol.RPC_RESPONSE, + payload=json.dumps({"dps": {"102": json.dumps({"id": 123, "result": [{"state": 8}]})}}).encode(), + ) + await peer.publish(response_topic, MessageParser.build(response, local_key=LOCAL_KEY, prefixed=False)) + received = await asyncio.wait_for(messages.get(), timeout=5) + assert received.protocol == response.protocol + assert received.payload == response.payload + unsub() diff --git a/uv.lock b/uv.lock index 8d13c80f..50712c23 100644 --- a/uv.lock +++ b/uv.lock @@ -1373,6 +1373,7 @@ dev = [ { name = "python-roborock", extra = ["cli"] }, { name = "ruff" }, { name = "syrupy" }, + { name = "zmqtt" }, ] [package.metadata] @@ -1409,6 +1410,7 @@ dev = [ { name = "python-roborock", extras = ["cli"] }, { name = "ruff", specifier = "==0.14.11" }, { name = "syrupy", specifier = ">=4.9.1,<6" }, + { name = "zmqtt", specifier = ">=0.2.0,<0.3.0" }, ] [[package]] @@ -1728,3 +1730,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/48/b7/503c98092fb3b344a179579f55814b613c1fbb1c23b3ec14a7b008a66a6e/yarl-1.22.0-cp314-cp314t-win_arm64.whl", hash = "sha256:9f6d73c1436b934e3f01df1e1b21ff765cd1d28c77dfb9ace207f746d4610ee1", size = 85171, upload-time = "2025-10-06T14:12:16.935Z" }, { url = "https://files.pythonhosted.org/packages/73/ae/b48f95715333080afb75a4504487cbe142cae1268afc482d06692d605ae6/yarl-1.22.0-py3-none-any.whl", hash = "sha256:1380560bdba02b6b6c90de54133c81c9f2a453dee9912fe58c1dcced1edb7cff", size = 46814, upload-time = "2025-10-06T14:12:53.872Z" }, ] + +[[package]] +name = "zmqtt" +version = "0.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ac/e6/23be7be9c899681c742bbbd861069f59e9cc669de929962951f23f869c6d/zmqtt-0.2.1.tar.gz", hash = "sha256:2b1ed7af2fb0e0eb0ee7d7a0477812e434d78e9626ebc0ef233690fc508118d8", size = 42724, upload-time = "2026-08-30T21:13:11.518Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/83/38/5e1a552960bfe5beabeb5fa32b78b71105a6289a11769c6772810f4eb181/zmqtt-0.2.1-py3-none-any.whl", hash = "sha256:af8e71066c55cc4802fd8352cb230133c7026314536fd92b62c3112848ecd8af", size = 54597, upload-time = "2026-08-30T21:13:10.202Z" }, +]