test: extract connected client to a fixture

style: fix formatting of tests

refactor: extract variables for mock data used in mqtt tests

style: fix lint errors in tests

Author: allenporter

tests/conftest.py

@@ -1,8 +1,8 @@
 import io
 import logging
-import queue
 import re
 from collections.abc import Callable, Generator
+from queue import Queue
 from typing import Any
 from unittest.mock import Mock, patch
 
@@ -63,29 +63,29 @@ def handle_socket_send(self, client_request: bytes) -> int:
 
 
 @pytest.fixture(name="received_requests")
-def received_requests_fixture() -> queue.Queue[bytes]:
+def received_requests_fixture() -> Queue[bytes]:
     """Fixture that provides access to the received requests."""
-    return queue.Queue()
+    return Queue()
 
 
 @pytest.fixture(name="response_queue")
-def response_queue_fixture() -> queue.Queue[bytes]:
+def response_queue_fixture() -> Generator[Queue[bytes], None, None]:
     """Fixture that provides access to the received requests."""
-    return queue.Queue()
+    response_queue: Queue[bytes] = Queue()
+    yield response_queue
+    assert response_queue.empty(), "Not all fake responses were consumed"
 
 
 @pytest.fixture(name="request_handler")
-def request_handler_fixture(
-    received_requests: queue.Queue[bytes], response_queue: queue.Queue[bytes]
-) -> RequestHandler:
+def request_handler_fixture(received_requests: Queue[bytes], response_queue: Queue[bytes]) -> RequestHandler:
     """Fixture records incoming requests and replies with responses from the queue."""
 
     def handle_request(client_request: bytes) -> bytes | None:
         """Handle an incoming request from the client."""
         received_requests.put(client_request)
 
         # Insert a prepared response into the response buffer
-        if response_queue.qsize() > 0:
+        if not response_queue.empty():
             return response_queue.get()
         return None
 

tests/mock_data.py

@@ -1,9 +1,14 @@
 """Mock data for Roborock tests."""
+
+import hashlib
+
 # All data is based on a U.S. customer with a Roborock S7 MaxV Ultra
 USER_EMAIL = "user@domain.com"
 
 BASE_URL = "https://usiot.roborock.com"
 
+USER_ID = "user123"
+K_VALUE = "domain123"
 USER_DATA = {
     "uid": 123456,
     "tokentype": "token_type",
@@ -14,10 +19,10 @@
     "country": "US",
     "nickname": "user_nickname",
     "rriot": {
-        "u": "user123",
+        "u": USER_ID,
         "s": "pass123",
         "h": "unknown123",
-        "k": "domain123",
+        "k": K_VALUE,
         "r": {
             "r": "US",
             "a": "https://api-us.roborock.com",
@@ -29,6 +34,7 @@
     "avatarurl": "https://files.roborock.com/iottest/default_avatar.png",
 }
 LOCAL_KEY = "key123"
+PRODUCT_ID = "product-id-123"
 HOME_DATA_RAW = {
     "id": 123456,
     "name": "My Home",
@@ -37,7 +43,7 @@
     "geoName": None,
     "products": [
         {
-            "id": "abc123",
+            "id": PRODUCT_ID,
             "name": "Roborock S7 MaxV",
             "code": "a27",
             "model": "roborock.vacuum.a27",
@@ -339,5 +345,5 @@
 }
 
 GET_CODE_RESPONSE = {"code": 200, "msg": "success", "data": None}
-
-MQTT_PUBLISH_TOPIC = "rr/m/o/user123/6ac2e6f8/abc123"
+HASHED_USER = hashlib.md5((USER_ID + ":" + K_VALUE).encode()).hexdigest()[2:10]
+MQTT_PUBLISH_TOPIC = f"rr/m/o/{USER_ID}/{HASHED_USER}/{PRODUCT_ID}"

tests/test_api.py

@@ -1,4 +1,5 @@
 import json
+from collections.abc import AsyncGenerator
 from queue import Queue
 from typing import Any
 from unittest.mock import patch
@@ -145,19 +146,28 @@ async def test_get_prop():
         assert props.dock_summary.dust_collection_mode is not None
 
 
-async def test_async_connect(
-    received_requests: Queue, response_queue: Queue, mqtt_client: RoborockMqttClientV1
-) -> None:
-    """Test connecting to the MQTT broker."""
-
+@pytest.fixture(name="connected_mqtt_client")
+async def connected_mqtt_client_fixture(
+    response_queue: Queue, mqtt_client: RoborockMqttClientV1
+) -> AsyncGenerator[RoborockMqttClientV1, None]:
     response_queue.put(mqtt_packet.gen_connack(rc=0, flags=2))
     response_queue.put(mqtt_packet.gen_suback(1, 0))
-
     await mqtt_client.async_connect()
-    assert mqtt_client.is_connected()
+    yield mqtt_client
+    if mqtt_client.is_connected():
+        await mqtt_client.async_disconnect()
 
-    await mqtt_client.async_disconnect()
-    assert not mqtt_client.is_connected()
+
+async def test_async_connect(received_requests: Queue, connected_mqtt_client: RoborockMqttClientV1) -> None:
+    """Test connecting to the MQTT broker."""
+
+    assert connected_mqtt_client.is_connected()
+    # Connecting again is a no-op
+    await connected_mqtt_client.async_connect()
+    assert connected_mqtt_client.is_connected()
+
+    await connected_mqtt_client.async_disconnect()
+    assert not connected_mqtt_client.is_connected()
 
     # Broker received a connect and subscribe. Disconnect packet is not
     # guaranteed to be captured by the time the async_disconnect returns
@@ -198,15 +208,10 @@ def build_rpc_response(message: dict[str, Any]) -> bytes:
 async def test_get_room_mapping(
     received_requests: Queue,
     response_queue: Queue,
-    mqtt_client: RoborockMqttClientV1,
+    connected_mqtt_client: RoborockMqttClientV1,
 ) -> None:
     """Test sending an arbitrary MQTT message and parsing the response."""
 
-    response_queue.put(mqtt_packet.gen_connack(rc=0, flags=2))
-    response_queue.put(mqtt_packet.gen_suback(1, 0))
-    await mqtt_client.async_connect()
-    assert mqtt_client.is_connected()
-
     test_request_id = 5050
     message = build_rpc_response(
         {
@@ -217,11 +222,9 @@ async def test_get_room_mapping(
     response_queue.put(mqtt_packet.gen_publish(MQTT_PUBLISH_TOPIC, payload=message))
 
     with patch("roborock.version_1_apis.roborock_client_v1.get_next_int", return_value=test_request_id):
-        room_mapping = await mqtt_client.get_room_mapping()
+        room_mapping = await connected_mqtt_client.get_room_mapping()
 
     assert room_mapping == [
         RoomMapping(segment_id=16, iot_id="2362048"),
         RoomMapping(segment_id=17, iot_id="2362044"),
     ]
-
-    await mqtt_client.async_disconnect()

tests/test_containers.py

@@ -10,7 +10,7 @@
     RoborockStateCode,
 )
 
-from .mock_data import CLEAN_RECORD, CLEAN_SUMMARY, CONSUMABLE, DND_TIMER, HOME_DATA_RAW, STATUS, USER_DATA
+from .mock_data import CLEAN_RECORD, CLEAN_SUMMARY, CONSUMABLE, DND_TIMER, HOME_DATA_RAW, PRODUCT_ID, STATUS, USER_DATA
 
 
 def test_user_data():
@@ -43,7 +43,7 @@ def test_home_data():
     assert hd.lat is None
     assert hd.geo_name is None
     product = hd.products[0]
-    assert product.id == "abc123"
+    assert product.id == PRODUCT_ID
     assert product.name == "Roborock S7 MaxV"
     assert product.code == "a27"
     assert product.model == "roborock.vacuum.a27"