client api coverage

Author: dmulcahey

tests/websocket/test_websocket_server_client.py

@@ -8,6 +8,7 @@
 from zha.application.gateway import WebSocketClientGateway, WebSocketServerGateway
 from zha.application.helpers import ZHAData
 from zha.websocket.client.client import Client
+from zha.websocket.client.helpers import ClientHelper
 
 
 async def test_server_client_connect_disconnect(
@@ -44,6 +45,36 @@ async def test_server_client_connect_disconnect(
     assert gateway._ws_server is None
 
 
+async def test_client_helper_disconnect(
+    zha_data: ZHAData,
+) -> None:
+    """Tests client helper disconnect logic."""
+
+    async with WebSocketServerGateway(zha_data) as gateway:
+        assert gateway.is_serving
+        assert gateway._ws_server is not None
+
+        client = Client(f"ws://localhost:{zha_data.ws_server_config.port}")
+        client_helper = ClientHelper(client)
+
+        await client.connect()
+        assert client.connected
+        assert "connected" in repr(client)
+
+        # The client does not begin listening immediately
+        assert client._listen_task is None
+        await client_helper.listen()
+        assert client._listen_task is not None
+
+        await client_helper.disconnect()
+        assert client._listen_task is None
+        assert "not connected" in repr(client)
+        assert not client.connected
+
+    assert not gateway.is_serving
+    assert gateway._ws_server is None
+
+
 @pytest.mark.parametrize(
     "zha_gateway",
     [

zha/application/gateway.py

@@ -969,15 +969,14 @@ async def connect(self) -> None:
             _LOGGER.exception("Unable to connect to the ZHA wss", exc_info=err)
             raise err
 
-        await self._client.listen()
-
     async def disconnect(self) -> None:
         """Disconnect from the websocket server."""
         await self._client.disconnect()
 
     async def __aenter__(self) -> WebSocketClientGateway:
         """Connect to the websocket server."""
         await self.connect()
+        await self.clients.listen()
         return self
 
     async def __aexit__(

zha/websocket/client/helpers.py

@@ -804,17 +804,19 @@ def __init__(self, client: Client):
     async def listen(self) -> WebSocketCommandResponse:
         """Listen for incoming messages."""
         command = ClientListenCommand()
+        await self._client.listen()
         return await self._client.async_send_command(command)
 
     async def listen_raw_zcl(self) -> WebSocketCommandResponse:
         """Listen for incoming raw ZCL messages."""
         command = ClientListenRawZCLCommand()
         return await self._client.async_send_command(command)
 
-    async def disconnect(self) -> WebSocketCommandResponse:
+    async def disconnect(self) -> None:
         """Disconnect this client from the server."""
         command = ClientDisconnectCommand()
-        return await self._client.async_send_command(command)
+        await self._client.async_send_command_no_wait(command)
+        await self._client.disconnect()
 
 
 class GroupHelper: