coverage

Author: dmulcahey

tests/conftest.py

@@ -418,6 +418,7 @@ async def __aenter__(self) -> CombinedWebsocketGateways:
         await client_gateway.clients.listen()
         await ws_gateway.async_block_till_done()
         await client_gateway.async_initialize()
+        assert client_gateway.state is not None
 
         self.combined_gateways = CombinedWebsocketGateways(
             self.zha_data, ws_gateway, client_gateway

tests/websocket/test_client_controller.py

@@ -178,6 +178,16 @@ async def test_ws_client_gateway_devices(
     await zha_gateway.async_block_till_done()
     assert len(ws_client_gateway.devices) == 2
 
+    # test client gateway device removal
+    await ws_client_gateway.async_remove_device(zha_device.ieee)
+    await zha_gateway.async_block_till_done()
+    assert len(ws_client_gateway.devices) == 1
+
+    # let's add it back
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_device)
+    await zha_gateway.async_block_till_done()
+    assert len(ws_client_gateway.devices) == 2
+
     # we removed and joined the device again so lets get the entity again
     client_device = ws_client_gateway.devices.get(zha_device.ieee)
     assert client_device is not None
@@ -398,6 +408,15 @@ async def test_ws_client_gateway_groups(
     assert client_device1.ieee in response.members_by_ieee
     assert client_device2.ieee in response.members_by_ieee
 
+    group_from_ws_client_gateway = ws_client_gateway.get_group(response.group_id)
+    assert group_from_ws_client_gateway is not None
+    assert group_from_ws_client_gateway.group_id == response.group_id
+    assert group_from_ws_client_gateway.name == response.name
+    assert (
+        group_from_ws_client_gateway.info_object.members_by_ieee
+        == response.members_by_ieee
+    )
+
     # test remove member from group from ws_client_gateway
     response = await ws_client_gateway.groups_helper.remove_group_members(
         response,

tests/websocket/test_websocket_server_client.py

@@ -2,11 +2,14 @@
 
 from __future__ import annotations
 
+from unittest.mock import patch
+
 import pytest
 
 from tests.conftest import CombinedWebsocketGateways
 from zha.application.gateway import WebSocketClientGateway, WebSocketServerGateway
 from zha.application.helpers import ZHAData
+from zha.websocket import ZHAWebSocketException
 from zha.websocket.client.client import Client
 from zha.websocket.client.helpers import ClientHelper
 
@@ -44,6 +47,13 @@ async def test_server_client_connect_disconnect(
     assert not gateway.is_serving
     assert gateway._ws_server is None
 
+    with (
+        pytest.raises(ZHAWebSocketException),
+        patch("zha.websocket.client.client.Client.connect", side_effect=TimeoutError),
+    ):
+        async with WebSocketClientGateway(zha_data) as client_gateway:
+            assert client_gateway.client.connected
+
 
 async def test_client_helper_disconnect(
     zha_data: ZHAData,

zha/application/gateway.py

@@ -78,6 +78,7 @@
     gather_with_limited_concurrency,
 )
 from zha.event import EventBase
+from zha.websocket import ZHAWebSocketException
 from zha.websocket.client.client import Client
 from zha.websocket.client.helpers import (
     AlarmControlPanelHelper,
@@ -113,10 +114,6 @@
 
 if TYPE_CHECKING:
     from zha.application.platforms.events import EntityStateChangedEvent
-    from zha.websocket.server.api.model import (
-        WebSocketCommand,
-        WebSocketCommandResponse,
-    )
     from zha.zigbee.model import ExtendedDeviceInfo, ZHAEvent
 
 BLOCK_LOG_TIMEOUT: Final[int] = 60
@@ -913,10 +910,10 @@ def __init__(self, config: ZHAData) -> None:
         self._client: Client = Client(
             self._ws_server_url, aiohttp_session=config.ws_client_config.aiohttp_session
         )
+        self._state: State
         self._devices: dict[EUI64, WebSocketClientDevice] = {}
         self._groups: dict[int, WebSocketClientGroup] = {}
         self.coordinator_zha_device: WebSocketClientDevice = None  # type: ignore[assignment]
-        self._state: State
         self.lights: LightHelper = LightHelper(self._client)
         self.switches: SwitchHelper = SwitchHelper(self._client)
         self.sirens: SirenHelper = SirenHelper(self._client)
@@ -965,9 +962,10 @@ async def connect(self) -> None:
         try:
             async with timeout(CONNECT_TIMEOUT):
                 await self._client.connect()
-        except Exception as err:
+        except TimeoutError as err:
             _LOGGER.exception("Unable to connect to the ZHA wss", exc_info=err)
-            raise err
+            await self._client.disconnect()
+            raise ZHAWebSocketException from err
 
     async def disconnect(self) -> None:
         """Disconnect from the websocket server."""
@@ -992,10 +990,6 @@ def create_and_track_task(self, coroutine: Coroutine) -> asyncio.Task:
         task.add_done_callback(self._tasks.remove)
         return task
 
-    async def send_command(self, command: WebSocketCommand) -> WebSocketCommandResponse:
-        """Send a command and get a response."""
-        return await self._client.async_send_command(command)
-
     async def load_devices(self) -> None:
         """Restore ZHA devices from zigpy application state."""
         response_devices = await self.devices_helper.get_devices()
@@ -1104,7 +1098,7 @@ def handle_state_changed(self, event: EntityStateChangedEvent) -> None:
     def handle_zha_event(self, event: ZHAEvent) -> None:
         """Handle a zha_event from the websocket server."""
         _LOGGER.debug("zha_event: %s", event)
-        device = self.devices.get(event.device.ieee)
+        device = self.devices.get(event.device_ieee)
         if device is None:
             _LOGGER.warning("Received zha_event from unknown device: %s", event)
             return

zha/websocket/client/client.py

@@ -163,8 +163,8 @@ async def disconnect(self) -> None:
 
             self._listen_task = None
 
-        assert self._client is not None
-        await self._client.close()
+        if self._client is not None:
+            await self._client.close()
 
         if self._close_aiohttp_session:
             await self.aiohttp_session.close()

zha/websocket/server/client.py

@@ -247,7 +247,6 @@ async def disconnect(
     gateway: WebSocketServerGateway, client: Client, command: WebSocketCommand
 ) -> None:
     """Disconnect the client."""
-    client.disconnect()
     gateway.client_manager.remove_client(client)
 
 

zha/zigbee/device.py

@@ -681,14 +681,18 @@ def update_available(
 
     def emit_zha_event(self, event_data: dict[str, str | int]) -> None:  # pylint: disable=unused-argument
         """Relay events directly."""
-        self.emit(
-            ZHA_EVENT,
-            ZHAEvent(
-                device_ieee=self.ieee,
-                unique_id=str(self.ieee),
-                data=event_data,
-            ),
+        event: ZHAEvent = ZHAEvent(
+            device_ieee=self.ieee,
+            unique_id=str(self.ieee),
+            data=event_data,
         )
+        self.emit(ZHA_EVENT, event)
+
+        # pylint: disable=import-outside-toplevel
+        from zha.application.gateway import WebSocketServerGateway
+
+        if isinstance(self.gateway, WebSocketServerGateway):
+            self.gateway.emit(ZHA_EVENT, event)
 
     async def _async_became_available(self) -> None:
         """Update device availability and signal entities."""