chore: Small tweaks to test fixtures

These are improvements factored out of a large change to add more e2e tests for device manager.

Author: allenporter

roborock/devices/local_channel.py

@@ -45,6 +45,11 @@ def connection_lost(self, exc: Exception | None) -> None:
         self.connection_lost_cb(exc)
 
 
+def get_running_loop() -> asyncio.AbstractEventLoop:
+    """Get the running event loop, extracted for mocking purposes."""
+    return asyncio.get_running_loop()
+
+
 class LocalChannel(Channel):
     """Simple RPC-style channel for communicating with a device over a local network.
 
@@ -179,7 +184,7 @@ async def connect(self) -> None:
         if self._is_connected:
             self._logger.debug("Unexpected call to connect when already connected")
             return
-        loop = asyncio.get_running_loop()
+        loop = get_running_loop()
         protocol = _LocalProtocol(self._data_received, self._connection_lost)
         try:
             self._transport, self._protocol = await loop.create_connection(lambda: protocol, self._host, _PORT)

tests/devices/test_local_channel.py

@@ -52,7 +52,7 @@ def setup_mock_loop(mock_transport: Mock) -> Generator[Mock, None, None]:
     loop = Mock()
     loop.create_connection = AsyncMock(return_value=(mock_transport, Mock()))
 
-    with patch("asyncio.get_running_loop", return_value=loop):
+    with patch("roborock.devices.local_channel.get_running_loop", return_value=loop):
         yield loop
 
 

tests/fixtures/aiomqtt_fixtures.py

@@ -28,10 +28,13 @@ async def mock_aiomqtt_client_fixture() -> AsyncGenerator[None, None]:
 
     async def poll_sockets(client: mqtt.Client) -> None:
         """Poll the mqtt client sockets in a loop to pick up new data."""
-        while True:
-            event_loop.call_soon_threadsafe(client.loop_read)
-            event_loop.call_soon_threadsafe(client.loop_write)
-            await asyncio.sleep(0.01)
+        try:
+            while True:
+                event_loop.call_soon_threadsafe(client.loop_read)
+                event_loop.call_soon_threadsafe(client.loop_write)
+                await asyncio.sleep(0.01)
+        except asyncio.CancelledError:
+            pass
 
     task: asyncio.Task[None] | None = None
 
@@ -52,6 +55,7 @@ def new_client(*args: Any, **kwargs: Any) -> mqtt.Client:
         yield
         if task:
             task.cancel()
+            await task
 
 
 @pytest.fixture

tests/fixtures/local_async_fixtures.py

@@ -79,7 +79,7 @@ def start_handle_write(data: bytes) -> None:
 
         return (mock_transport, protocol)
 
-    with patch("roborock.devices.local_channel.asyncio.get_running_loop") as mock_loop:
+    with patch("roborock.devices.local_channel.get_running_loop") as mock_loop:
         mock_loop.return_value.create_connection.side_effect = create_connection
         yield
 

tests/fixtures/mqtt.py

@@ -32,6 +32,7 @@ def __init__(
         self.handle_request = handle_request
         self.response_queue = response_queue
         self.log = log
+        self.client_connected = False
 
     def pending(self) -> int:
         """Return the number of bytes in the response buffer."""
@@ -45,31 +46,31 @@ def handle_socket_recv(self, read_size: int) -> bytes:
         self.response_buf.seek(0)
         data = self.response_buf.read(read_size)
         _LOGGER.debug("Response: 0x%s", data.hex())
+        self.log.add_log_entry("[mqtt <]", data)
         # Consume the rest of the data in the buffer
         remaining_data = self.response_buf.read()
         self.response_buf = io.BytesIO(remaining_data)
         return data
 
     def handle_socket_send(self, client_request: bytes) -> int:
         """Receive an incoming request from the client."""
+        self.client_connected = True
         _LOGGER.debug("Request: 0x%s", client_request.hex())
         self.log.add_log_entry("[mqtt >]", client_request)
         if (response := self.handle_request(client_request)) is not None:
             # Enqueue a response to be sent back to the client in the buffer.
             # The buffer will be emptied when the client calls recv() on the socket
             _LOGGER.debug("Queued: 0x%s", response.hex())
-            self.log.add_log_entry("[mqtt <]", response)
             self.response_buf.write(response)
         return len(client_request)
 
     def push_response(self) -> None:
         """Push a response to the client."""
-        if not self.response_queue.empty():
+        if not self.response_queue.empty() and self.client_connected:
             response = self.response_queue.get()
             # Enqueue a response to be sent back to the client in the buffer.
             # The buffer will be emptied when the client calls recv() on the socket
             _LOGGER.debug("Queued: 0x%s", response.hex())
-            self.log.add_log_entry("[mqtt <]", response)
             self.response_buf.write(response)