test: Extract connected client to a fixture

Author: allenporter

tests/conftest.py

@@ -1,6 +1,6 @@
 import io
 import logging
-import queue
+from queue import Queue
 import re
 from collections.abc import Callable, Generator
 from typing import Any
@@ -63,20 +63,22 @@ 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() -> Queue[bytes]:
     """Fixture that provides access to the received requests."""
-    return queue.Queue()
+    response_queue = 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]
+    received_requests: Queue[bytes], response_queue: Queue[bytes]
 ) -> RequestHandler:
     """Fixture records incoming requests and replies with responses from the queue."""
 
@@ -85,7 +87,7 @@ def handle_request(client_request: bytes) -> bytes | None:
         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/test_api.py

@@ -2,6 +2,7 @@
 from queue import Queue
 from typing import Any
 from unittest.mock import patch
+from collections.abc import AsyncGenerator
 
 import paho.mqtt.client as mqtt
 import pytest
@@ -145,19 +146,28 @@ async def test_get_prop():
         assert props.dock_summary.dust_collection_mode is not None
 
 
+@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()
+    yield mqtt_client
+    if mqtt_client.is_connected():
+        await mqtt_client.async_disconnect()
+
+
 async def test_async_connect(
-    received_requests: Queue, response_queue: Queue, mqtt_client: RoborockMqttClientV1
+    received_requests: Queue, connected_mqtt_client: RoborockMqttClientV1
 ) -> None:
     """Test connecting to the MQTT broker."""
 
-    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()
+    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 mqtt_client.async_disconnect()
-    assert not 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()