diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d19150576..4ce647d01 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -55,7 +55,38 @@ jobs: token: ${{ secrets.CODECOV_TOKEN }} slug: Python-roborock/python-roborock - + test-mqtt: + strategy: + fail-fast: false + matrix: + python-version: + - "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 + uses: astral-sh/setup-uv@v7 + with: + python-version: ${{ matrix.python-version }} + activate-environment: true + - name: Test MQTT with Pytest + run: uv run pytest -m mqtt_broker --log-cli-level=DEBUG -vv -s --cov --cov-branch --cov-report=xml + - name: Upload results to Codecov + uses: codecov/codecov-action@v5 + with: + token: ${{ secrets.CODECOV_TOKEN }} + slug: Python-roborock/python-roborock build: name: Build Package @@ -110,6 +141,7 @@ jobs: runs-on: ubuntu-latest needs: - test + - test-mqtt concurrency: release if: github.ref == 'refs/heads/main' permissions: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8776efe13..6872fae84 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -55,11 +55,25 @@ pre-commit run --all-files We use `pytest` for testing. Please ensure all tests pass and add new tests for your changes. +Run the standard test suite without Docker: + +```bash +uv run pytest +``` + +Real-broker MQTT tests are excluded by default. To run them, start EMQX at +`127.0.0.1:1888` with Docker Compose and select the `mqtt_broker` marker: + ```bash -# Run tests -pytest +docker compose up -d --wait +uv run pytest -m mqtt_broker +docker compose down ``` +CI runs both suites in separate parallel jobs on Python 3.11 and 3.14, with an +EMQX service container only in the MQTT job. Each test has a 30-second timeout; +MQTT connection and message waits have shorter timeouts. + ## Pull Requests 1. **Create a branch** for your changes. diff --git a/compose.yaml b/compose.yaml new file mode 100644 index 000000000..1e186fa5b --- /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 7a6a9ffdc..fc031607a 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] @@ -146,6 +147,8 @@ module = ["roborock.map.proto.*"] ignore_errors = true [tool.pytest.ini_options] +addopts = "-m 'not mqtt_broker'" +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 000000000..1ce3ee933 --- /dev/null +++ b/tests/e2e/test_mqtt_broker.py @@ -0,0 +1,249 @@ +"""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 asyncio.wait_for(create_mqtt_session(mqtt_params), timeout=mqtt_params.timeout) + try: + assert session.connected + yield session + finally: + await asyncio.wait_for(session.close(), timeout=mqtt_params.timeout) + + +@pytest.fixture(name="peer") +async def peer_fixture(mqtt_params: MqttParams) -> AsyncGenerator[MQTTClientV5, None]: + """Represent a device using an independent MQTT client.""" + peer = create_client( + mqtt_params.host, + mqtt_params.port, + version="5.0", + mqtt_connect_timeout=mqtt_params.timeout, + reconnect=ReconnectConfig(enabled=False), + ) + try: + await asyncio.wait_for(peer.connect(), timeout=mqtt_params.timeout) + yield peer + finally: + await asyncio.wait_for(peer.disconnect(), timeout=mqtt_params.timeout) + + +@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 asyncio.wait_for(peer.publish(topic, payload), timeout=5) + 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 asyncio.timeout(5), 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 asyncio.wait_for(peer.publish(topic, b"both"), timeout=5) + 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 asyncio.wait_for(peer.publish(topic, b"second only"), timeout=5) + 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 asyncio.wait_for(peer.publish(topic, b"first again"), timeout=5) + 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 asyncio.wait_for(peer.publish(f"{topic}/first", b"first"), timeout=5) + await asyncio.wait_for(peer.publish(f"{topic}/second", b"second"), timeout=5) + + 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, mqtt_params: MqttParams, 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 asyncio.wait_for(peer.publish(topic, b"before restart"), timeout=5) + assert await asyncio.wait_for(messages.get(), timeout=5) == b"before restart" + + completed_connections = mqtt_params.diagnostics.as_dict().get("connection_count", 0) + await session.restart() + async with asyncio.timeout(20): + # The counter records the disconnect even if reconnection finishes between polls. + while ( + mqtt_params.diagnostics.as_dict().get("connection_count", 0) <= completed_connections + or not session.connected + ): + await asyncio.sleep(0.01) + + await asyncio.wait_for(peer.publish(topic, b"after restart"), timeout=5) + assert await asyncio.wait_for(messages.get(), timeout=5) == b"after restart" + async with asyncio.timeout(5), 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 asyncio.wait_for(peer.publish(topic, b"inbound"), timeout=5) + 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 asyncio.timeout(5), 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 asyncio.timeout(5), 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 8448eb46c..929c44208 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" }, +]