Fix ws receive task (#11)

Author: edenhaus

go2rtc_client/ws/client.py

@@ -76,6 +76,12 @@ async def close(self) -> None:
             self._client = None
             await client.close()
 
+        if self._rx_task:
+            task = self._rx_task
+            self._rx_task = None
+            task.cancel()
+            await task
+
     @handle_error
     async def send(self, message: SendMessages) -> None:
         """Send a message."""
@@ -110,34 +116,23 @@ async def _receive_messages(self) -> None:
         if TYPE_CHECKING:
             assert self._client
 
-        try:
-            while self.connected:
-                msg = await self._client.receive()
-                match msg.type:
-                    case (
-                        WSMsgType.CLOSE
-                        | WSMsgType.CLOSED
-                        | WSMsgType.CLOSING
-                        | WSMsgType.PING
-                        | WSMsgType.PONG
-                    ):
-                        break
-                    case WSMsgType.ERROR:
-                        _LOGGER.error("Error received: %s", msg.data)
-                    case WSMsgType.TEXT:
-                        self._process_text_message(msg.data)
-                    case _:
-                        _LOGGER.warning("Received unknown message: %s", msg)
-        except Exception:
-            _LOGGER.exception("Unexpected error while receiving message")
-            raise
-        finally:
-            _LOGGER.debug(
-                "Websocket client connection from %s closed", self._server_url
-            )
-
-            if self.connected:
-                await self.close()
+        while self.connected:
+            msg = await self._client.receive()
+            match msg.type:
+                case (
+                    WSMsgType.CLOSE
+                    | WSMsgType.CLOSED
+                    | WSMsgType.CLOSING
+                    | WSMsgType.PING
+                    | WSMsgType.PONG
+                ):
+                    break
+                case WSMsgType.ERROR:
+                    _LOGGER.error("Error received: %s", msg.data)
+                case WSMsgType.TEXT:
+                    self._process_text_message(msg.data)
+                case _:
+                    _LOGGER.warning("Received unknown message: %s", msg)
 
     def subscribe(
         self, callback: Callable[[ReceiveMessages], None]

pyproject.toml

@@ -56,7 +56,7 @@ include = [
 
 [tool.coverage.report]
 show_missing = true
-fail_under = 50
+fail_under = 97
 
 [tool.coverage.run]
 plugins = ["covdefaults"]

tests/ws/test_client.py

@@ -34,6 +34,8 @@
 class TestServer:
     """Test server."""
 
+    __test__ = False
+
     def __init__(self) -> None:
         """Initialize the test server."""
         self.server: AioHttpTestServer
@@ -361,32 +363,3 @@ async def receive() -> WSMessage:
     await asyncio.sleep(0.1)
 
     assert caplog.record_tuples == [record]
-
-
-async def test_receive_raised(
-    caplog: pytest.LogCaptureFixture,
-    ws_client: Go2RtcWsClient,
-) -> None:
-    """Test getting message raised an exception."""
-    client = AsyncMock()
-    client.return_value.closed = False
-    ws_client._session.ws_connect = client  # type: ignore[method-assign] # pylint: disable=protected-access
-
-    async def receive() -> WSMessage:
-        nonlocal client
-        client.return_value.closed = True
-
-        raise ValueError
-
-    client.return_value.receive.side_effect = receive
-
-    await ws_client.connect()
-    await asyncio.sleep(0.1)
-
-    assert caplog.record_tuples == [
-        (
-            "go2rtc_client.ws.client",
-            logging.ERROR,
-            "Unexpected error while receiving message",
-        )
-    ]