Organize test fixtures to prepare for more end to end tests (#699)

* chore: Organize test fixtures

Rewrite the test fixtures to have a more clear split between local and mqtt fixtures. This is in prepration for running both at once for e2 tests for device manager.

All fixtures are moved into a fixtures subdirectory. The helper classes that are imported into other tests are added in separate files for importing, and to avoid import warnings from pytests.

This renames all the fixtures to have mqtt prefixed names and local fixtures to have local prefixed names. There is one minor change to make the local asyncio tests uses asyncio Queues rather than blocking queues.

* chore: Address co-pilot review feedback

* chore: fix lint errors

* chore: Remove duplicate captured request log

Author: allenporter

tests/conftest.py

@@ -11,8 +11,7 @@
     RoborockMessage,
     RoborockMessageProtocol,
 )
-
-from ..conftest import FakeChannel
+from tests.fixtures.channel_fixtures import FakeChannel
 
 
 @pytest.fixture

tests/devices/test_a01_channel.py

@@ -25,9 +25,8 @@
 from roborock.protocols.v1_protocol import MapResponse, SecurityData, V1RpcChannel
 from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
 from roborock.roborock_typing import RoborockCommand
-
-from .. import mock_data
-from ..conftest import FakeChannel
+from tests import mock_data
+from tests.fixtures.channel_fixtures import FakeChannel
 
 USER_DATA = UserData.from_dict(mock_data.USER_DATA)
 TEST_DEVICE_UID = "abc123"

tests/devices/test_v1_channel.py

@@ -8,7 +8,7 @@
 
 from roborock.devices.traits.a01 import DyadApi, ZeoApi
 from roborock.roborock_message import RoborockDyadDataProtocol, RoborockMessageProtocol, RoborockZeoProtocol
-from tests.conftest import FakeChannel
+from tests.fixtures.channel_fixtures import FakeChannel
 from tests.protocols.common import build_a01_message
 
 

tests/devices/traits/a01/test_init.py

@@ -18,7 +18,7 @@
 from roborock.exceptions import RoborockException
 from roborock.protocols.b01_protocol import B01_VERSION
 from roborock.roborock_message import RoborockB01Props, RoborockMessage, RoborockMessageProtocol
-from tests.conftest import FakeChannel
+from tests.fixtures.channel_fixtures import FakeChannel
 
 
 def build_b01_message(message: dict[Any, Any], msg_id: str = "123456789", seq: int = 2020) -> RoborockMessage:

tests/devices/traits/b01/test_init.py

@@ -1 +1,8 @@
 """End-to-end tests package."""
+
+pytest_plugins = [
+    "tests.fixtures.logging_fixtures",
+    "tests.fixtures.local_async_fixtures",
+    "tests.fixtures.pahomqtt_fixtures",
+    "tests.fixtures.aiomqtt_fixtures",
+]

tests/e2e/__init__.py

@@ -1,17 +1,14 @@
 """End-to-end tests for LocalChannel using fake sockets."""
 
 import asyncio
-from collections.abc import AsyncGenerator, Callable, Generator
-from queue import Queue
-from typing import Any
-from unittest.mock import Mock, patch
+from collections.abc import AsyncGenerator
+from unittest.mock import patch
 
 import pytest
 
 from roborock.devices.local_channel import LocalChannel
 from roborock.protocol import create_local_decoder, create_local_encoder
 from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
-from tests.conftest import RequestHandler
 from tests.mock_data import LOCAL_KEY
 
 TEST_HOST = "192.168.1.100"
@@ -21,35 +18,8 @@
 TEST_RANDOM = 13579
 
 
-@pytest.fixture(name="mock_create_local_connection")
-def create_local_connection_fixture(request_handler: RequestHandler) -> Generator[None, None, None]:
-    """Fixture that overrides the transport creation to wire it up to the mock socket."""
-
-    async def create_connection(protocol_factory: Callable[[], asyncio.Protocol], *args, **kwargs) -> tuple[Any, Any]:
-        protocol = protocol_factory()
-
-        def handle_write(data: bytes) -> None:
-            response = request_handler(data)
-            if response is not None:
-                # Call data_received directly to avoid loop scheduling issues in test
-                protocol.data_received(response)
-
-        closed = asyncio.Event()
-
-        mock_transport = Mock()
-        mock_transport.write = handle_write
-        mock_transport.close = closed.set
-        mock_transport.is_reading = lambda: not closed.is_set()
-
-        return (mock_transport, protocol)
-
-    with patch("roborock.devices.local_channel.asyncio.get_running_loop") as mock_loop:
-        mock_loop.return_value.create_connection.side_effect = create_connection
-        yield
-
-
 @pytest.fixture(name="local_channel")
-async def local_channel_fixture(mock_create_local_connection: None) -> AsyncGenerator[LocalChannel, None]:
+async def local_channel_fixture(mock_async_create_local_connection: None) -> AsyncGenerator[LocalChannel, None]:
     with patch(
         "roborock.devices.local_channel.get_next_int", return_value=TEST_CONNECT_NONCE, device_uid=TEST_DEVICE_UID
     ):
@@ -80,19 +50,23 @@ def build_response(
 
 
 async def test_connect(
-    local_channel: LocalChannel, response_queue: Queue[bytes], received_requests: Queue[bytes]
+    local_channel: LocalChannel,
+    local_response_queue: asyncio.Queue[bytes],
+    local_received_requests: asyncio.Queue[bytes],
 ) -> None:
     """Test connecting to the device."""
     # Queue HELLO response with payload to ensure it can be parsed
-    response_queue.put(build_response(RoborockMessageProtocol.HELLO_RESPONSE, 1, payload=b"ok", random=TEST_RANDOM))
+    local_response_queue.put_nowait(
+        build_response(RoborockMessageProtocol.HELLO_RESPONSE, 1, payload=b"ok", random=TEST_RANDOM)
+    )
 
     await local_channel.connect()
 
     assert local_channel.is_connected
-    assert received_requests.qsize() == 1
+    assert local_received_requests.qsize() == 1
 
     # Verify HELLO request
-    request_bytes = received_requests.get()
+    request_bytes = await local_received_requests.get()
     # Note: We cannot use create_local_decoder here because HELLO_REQUEST has payload=None
     # which causes MessageParser to fail parsing. For now we verify the raw bytes.
 
@@ -104,17 +78,21 @@ async def test_connect(
 
 
 async def test_send_command(
-    local_channel: LocalChannel, response_queue: Queue[bytes], received_requests: Queue[bytes]
+    local_channel: LocalChannel,
+    local_response_queue: asyncio.Queue[bytes],
+    local_received_requests: asyncio.Queue[bytes],
 ) -> None:
     """Test sending a command."""
     # Queue HELLO response
-    response_queue.put(build_response(RoborockMessageProtocol.HELLO_RESPONSE, 1, payload=b"ok", random=TEST_RANDOM))
+    local_response_queue.put_nowait(
+        build_response(RoborockMessageProtocol.HELLO_RESPONSE, 1, payload=b"ok", random=TEST_RANDOM)
+    )
 
     await local_channel.connect()
 
     # Clear requests from handshake
-    while not received_requests.empty():
-        received_requests.get()
+    while not local_received_requests.empty():
+        await local_received_requests.get()
 
     # Send command
     cmd_seq = 123
@@ -127,8 +105,8 @@ async def test_send_command(
     await local_channel.publish(msg)
 
     # Verify request
-    assert received_requests.qsize() == 1
-    request_bytes = received_requests.get()
+    request_bytes = await local_received_requests.get()
+    assert local_received_requests.empty()
 
     # Decode request
     decoder = create_local_decoder(local_key=LOCAL_KEY, connect_nonce=TEST_CONNECT_NONCE, ack_nonce=TEST_ACK_NONCE)

tests/e2e/test_local_session.py

@@ -18,12 +18,12 @@
 from roborock.protocol import MessageParser
 from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
 from tests import mqtt_packet
+from tests.fixtures.mqtt import FAKE_PARAMS, Subscriber
 from tests.mock_data import LOCAL_KEY
-from tests.mqtt_fixtures import FAKE_PARAMS, Subscriber
 
 
 @pytest.fixture(autouse=True)
-def auto_mock_mqtt_client(mock_mqtt_client_fixture: None) -> None:
+def auto_mock_mqtt_client(mock_aiomqtt_client: None) -> None:
     """Automatically use the mock mqtt client fixture."""
 
 
@@ -33,7 +33,7 @@ def auto_fast_backoff(fast_backoff_fixture: None) -> None:
 
 
 @pytest.fixture(autouse=True)
-def mqtt_server_fixture(mock_create_connection: None, mock_select: None) -> None:
+def mqtt_server_fixture(mock_paho_mqtt_create_connection: None, mock_paho_mqtt_select: None) -> None:
     """Fixture to mock the MQTT connection.
 
     This is here to pull in the mock socket pixtures into all tests used here.
@@ -42,10 +42,10 @@ def mqtt_server_fixture(mock_create_connection: None, mock_select: None) -> None
 
 @pytest.fixture(name="session")
 async def session_fixture(
-    push_response: Callable[[bytes], None],
+    push_mqtt_response: Callable[[bytes], None],
 ) -> AsyncGenerator[MqttSession, None]:
     """Fixture to create a new connected MQTT session."""
-    push_response(mqtt_packet.gen_connack(rc=0, flags=2))
+    push_mqtt_response(mqtt_packet.gen_connack(rc=0, flags=2))
     session = await create_mqtt_session(FAKE_PARAMS)
     assert session.connected
     try:
@@ -54,12 +54,12 @@ async def session_fixture(
         await session.close()
 
 
-async def test_session_e2e_receive_message(push_response: Callable[[bytes], None], session: MqttSession) -> None:
+async def test_session_e2e_receive_message(push_mqtt_response: Callable[[bytes], None], session: MqttSession) -> None:
     """Test receiving a real Roborock message through the session."""
     assert session.connected
 
     # Subscribe to the topic. We'll next construct and push a message.
-    push_response(mqtt_packet.gen_suback(mid=1))
+    push_mqtt_response(mqtt_packet.gen_suback(mid=1))
     subscriber = Subscriber()
     await session.subscribe("topic-1", subscriber.append)
 
@@ -71,12 +71,13 @@ async def test_session_e2e_receive_message(push_response: Callable[[bytes], None
     payload = MessageParser.build(msg, local_key=LOCAL_KEY, prefixed=False)
 
     # Simulate receiving the message from the broker
-    push_response(mqtt_packet.gen_publish("topic-1", mid=2, payload=payload))
+    push_mqtt_response(mqtt_packet.gen_publish("topic-1", mid=2, payload=payload))
 
     # Verify it was dispatched to the subscriber
     await subscriber.wait()
     assert len(subscriber.messages) == 1
     received_payload = subscriber.messages[0]
+    assert isinstance(received_payload, bytes)
     assert received_payload == payload
 
     # Verify the message payload contents
@@ -90,8 +91,8 @@ async def test_session_e2e_receive_message(push_response: Callable[[bytes], None
 
 
 async def test_session_e2e_publish_message(
-    push_response: Callable[[bytes], None],
-    received_requests: Queue,
+    push_mqtt_response: Callable[[bytes], None],
+    mqtt_received_requests: Queue,
     session: MqttSession,
 ) -> None:
     """Test publishing a real Roborock message."""
@@ -109,8 +110,8 @@ async def test_session_e2e_publish_message(
     # Verify what was sent to the broker
     # We expect the payload to be present in the sent bytes
     found = False
-    while not received_requests.empty():
-        request = received_requests.get()
+    while not mqtt_received_requests.empty():
+        request = mqtt_received_requests.get()
         if payload in request:
             found = True
             break

tests/e2e/test_mqtt_session.py

@@ -10,40 +10,11 @@
 import paho.mqtt.client as mqtt
 import pytest
 
-from roborock.mqtt.session import MqttParams
-from tests.conftest import FakeSocketHandler
+from .mqtt import FakeMqttSocketHandler
 
-FAKE_PARAMS = MqttParams(
-    host="localhost",
-    port=1883,
-    tls=False,
-    username="username",
-    password="password",
-    timeout=10.0,
-)
 
-
-class Subscriber:
-    """Mock subscriber class.
-
-    We use this to hold on to received messages for verification.
-    """
-
-    def __init__(self) -> None:
-        self.messages: list[bytes] = []
-        self._event = asyncio.Event()
-
-    def append(self, message: bytes) -> None:
-        self.messages.append(message)
-        self._event.set()
-
-    async def wait(self) -> None:
-        await asyncio.wait_for(self._event.wait(), timeout=1.0)
-        self._event.clear()
-
-
-@pytest.fixture
-async def mock_mqtt_client_fixture() -> AsyncGenerator[None, None]:
+@pytest.fixture(name="mock_aiomqtt_client")
+async def mock_aiomqtt_client_fixture() -> AsyncGenerator[None, None]:
     """Fixture to patch the MQTT underlying sync client.
 
     The tests use fake sockets, so this ensures that the async mqtt client does not
@@ -94,11 +65,13 @@ def fast_backoff_fixture() -> Generator[None, None, None]:
 
 
 @pytest.fixture
-def push_response(response_queue: Queue, fake_socket_handler: FakeSocketHandler) -> Callable[[bytes], None]:
+def push_mqtt_response(
+    mqtt_response_queue: Queue, fake_mqtt_socket_handler: FakeMqttSocketHandler
+) -> Callable[[bytes], None]:
     """Fixture to push a response to the client."""
 
     def _push(data: bytes) -> None:
-        response_queue.put(data)
-        fake_socket_handler.push_response()
+        mqtt_response_queue.put(data)
+        fake_mqtt_socket_handler.push_response()
 
     return _push