From 40d586f8ce10cf0e66fde4997616b19639cb344b Mon Sep 17 00:00:00 2001
From: "David F. Mulcahey" <david.mulcahey@me.com>
Date: Sun, 13 Feb 2022 12:08:52 -0500
Subject: [PATCH] ZHA events and device triggers (#33)

* start implementing zha event

* remove stale todo

* wrong word

* add unique_id

* add device automation triggers

* proxy zha events to devices
---
 zhaws/client/controller.py              | 11 ++++++++++
 zhaws/client/model/events.py            | 15 +++++++++++++-
 zhaws/client/model/types.py             |  2 ++
 zhaws/client/proxy.py                   | 11 +++++++++-
 zhaws/server/const.py                   |  8 ++++++++
 zhaws/server/platforms/__init__.py      |  1 +
 zhaws/server/platforms/discovery.py     |  2 +-
 zhaws/server/zigbee/cluster/__init__.py | 27 ++++++++++---------------
 zhaws/server/zigbee/cluster/general.py  |  2 +-
 zhaws/server/zigbee/cluster/security.py |  5 +----
 zhaws/server/zigbee/device.py           | 16 +++++++--------
 zhaws/server/zigbee/endpoint.py         | 14 +------------
 12 files changed, 68 insertions(+), 46 deletions(-)

diff --git a/zhaws/client/controller.py b/zhaws/client/controller.py
index 1562b1b7..a57755dd 100644
--- a/zhaws/client/controller.py
+++ b/zhaws/client/controller.py
@@ -41,6 +41,7 @@
     GroupRemovedEvent,
     PlatformEntityStateChangedEvent,
     RawDeviceInitializedEvent,
+    ZHAEvent,
 )
 from zhaws.client.proxy import DeviceProxy, GroupProxy
 from zhaws.event import EventBase
@@ -88,6 +89,7 @@ def __init__(
         self._client.on_event(
             EventTypes.PLATFORM_ENTITY_EVENT, self._handle_event_protocol
         )
+        self._client.on_event(EventTypes.DEVICE_EVENT, self._handle_event_protocol)
         self._client.on_event(EventTypes.CONTROLLER_EVENT, self._handle_event_protocol)
 
     @property
@@ -160,6 +162,15 @@ def handle_platform_entity_state_changed(
                 return
             group.emit_platform_entity_event(event)
 
+    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)
+        if device is None:
+            _LOGGER.warning("Received zha_event from unknown device: %s", event)
+            return
+        device.emit("zha_event", event)
+
     def handle_device_joined(self, event: DeviceJoinedEvent) -> None:
         """Handle device joined.
 
diff --git a/zhaws/client/model/events.py b/zhaws/client/model/events.py
index 022d87ed..ab4cfff6 100644
--- a/zhaws/client/model/events.py
+++ b/zhaws/client/model/events.py
@@ -41,6 +41,7 @@ class MinimalEndpoint(BaseModel):
     """Minimal endpoint model."""
 
     id: int
+    unique_id: str
 
 
 class MinimalDevice(BaseModel):
@@ -82,7 +83,6 @@ class MinimalGroup(BaseModel):
 class PlatformEntityStateChangedEvent(BaseEvent):
     """Platform entity event."""
 
-    """TODO use this as a base and create specific events for each entity type where state and attributes is fully modeled out"""
     event_type: Literal["platform_entity_event"] = "platform_entity_event"
     event: Literal["platform_entity_state_changed"] = "platform_entity_state_changed"
     platform_entity: MinimalPlatformEntity
@@ -196,6 +196,18 @@ class DeviceOnlineEvent(BaseEvent):
     device: MinimalDevice
 
 
+class ZHAEvent(BaseEvent):
+    """ZHA event."""
+
+    event: Literal["zha_event"] = "zha_event"
+    event_type: Literal["device_event"] = "device_event"
+    device: MinimalDevice
+    cluster_handler: MinimalClusterHandler
+    endpoint: MinimalEndpoint
+    command: str
+    args: Union[list, dict]
+
+
 class GroupRemovedEvent(ControllerEvent):
     """Group removed event."""
 
@@ -240,6 +252,7 @@ class GroupMemberRemovedEvent(ControllerEvent):
         GroupMemberRemovedEvent,
         DeviceOfflineEvent,
         DeviceOnlineEvent,
+        ZHAEvent,
     ],
     Field(discriminator="event"),  # noqa: F821
 ]
diff --git a/zhaws/client/model/types.py b/zhaws/client/model/types.py
index 0ffc4c29..108a0873 100644
--- a/zhaws/client/model/types.py
+++ b/zhaws/client/model/types.py
@@ -43,6 +43,7 @@ class Endpoint(BaseModel):
     """Endpoint model."""
 
     id: int
+    unique_id: str
 
 
 class GenericState(BaseModel):
@@ -464,6 +465,7 @@ class Device(BaseDevice):
         ],
     ]
     neighbors: list[Any]
+    device_automation_triggers: dict[str, dict[str, Any]]
 
 
 class GroupEntity(BaseEntity):
diff --git a/zhaws/client/proxy.py b/zhaws/client/proxy.py
index 12c7982e..79db3e07 100644
--- a/zhaws/client/proxy.py
+++ b/zhaws/client/proxy.py
@@ -1,7 +1,7 @@
 """Proxy object for the client side objects."""
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from zhaws.client.model.events import PlatformEntityStateChangedEvent
 from zhaws.client.model.types import ButtonEntity
@@ -87,5 +87,14 @@ def device_model(self, device_model: DeviceModel) -> None:
         """Set the device model."""
         self._proxied_object = device_model
 
+    @property
+    def device_automation_triggers(self) -> dict[tuple[str, str], dict[str, Any]]:
+        """Return the device automation triggers."""
+        model_triggers = self._proxied_object.device_automation_triggers
+        return {
+            (key.split("~")[0], key.split("~")[1]): value
+            for key, value in model_triggers.items()
+        }
+
     def __repr__(self) -> str:
         return self._proxied_object.__repr__()
diff --git a/zhaws/server/const.py b/zhaws/server/const.py
index d766eafc..78196061 100644
--- a/zhaws/server/const.py
+++ b/zhaws/server/const.py
@@ -128,6 +128,14 @@ class RawZCLEvents(StrEnum):
     ATTRIBUTE_UPDATED = "attribute_updated"
 
 
+class DeviceEvents(StrEnum):
+    """Events that devices can broadcast."""
+
+    DEVICE_OFFLINE = "device_offline"
+    DEVICE_ONLINE = "device_online"
+    ZHA_EVENT = "zha_event"
+
+
 ATTR_UNIQUE_ID: Final[str] = "unique_id"
 COMMAND: Final[str] = "command"
 CONF_BAUDRATE: Final[str] = "baudrate"
diff --git a/zhaws/server/platforms/__init__.py b/zhaws/server/platforms/__init__.py
index ceed4d38..9a9c7bd5 100644
--- a/zhaws/server/platforms/__init__.py
+++ b/zhaws/server/platforms/__init__.py
@@ -197,6 +197,7 @@ def send_event(self, signal: dict[str, Any]) -> None:
         }
         signal["endpoint"] = {
             "id": self._endpoint.id,
+            "unique_id": self._endpoint.unique_id,
         }
         _LOGGER.info("Sending event from platform entity: %s", signal)
         self.device.send_event(signal)
diff --git a/zhaws/server/platforms/discovery.py b/zhaws/server/platforms/discovery.py
index fec8e39b..106b1721 100644
--- a/zhaws/server/platforms/discovery.py
+++ b/zhaws/server/platforms/discovery.py
@@ -93,7 +93,7 @@ def __init__(self) -> None:
     def discover_entities(self, endpoint: Endpoint) -> None:
         """Process an endpoint on a zigpy device."""
         _LOGGER.info(
-            "Discovering entitied for endpoint: %s-%s",
+            "Discovering entities for endpoint: %s-%s",
             str(endpoint.device.ieee),
             endpoint.id,
         )
diff --git a/zhaws/server/zigbee/cluster/__init__.py b/zhaws/server/zigbee/cluster/__init__.py
index 26b6b5aa..f5502559 100644
--- a/zhaws/server/zigbee/cluster/__init__.py
+++ b/zhaws/server/zigbee/cluster/__init__.py
@@ -14,7 +14,7 @@
 
 from zhaws.event import EventBase
 from zhaws.model import BaseEvent
-from zhaws.server.const import EVENT, EVENT_TYPE, EventTypes
+from zhaws.server.const import EVENT, EVENT_TYPE, DeviceEvents, EventTypes
 from zhaws.server.util import LogMixin
 from zhaws.server.zigbee.cluster.const import (
     CLUSTER_HANDLER_ZDO,
@@ -350,18 +350,16 @@ def attribute_updated(self, attrid: int, value: Any) -> None:
     def zdo_command(self, *args: Any, **kwargs: Any) -> None:
         """Handle ZDO commands on this cluster."""
 
-    def zha_send_event(self, command: str, args: int | dict) -> None:
+    def zha_send_event(self, command: str, args: list | dict) -> None:
         """Relay events to hass."""
-        """ TODO
-        self._ch_pool.zha_send_event(
+        self.send_event(
             {
-                ATTR_UNIQUE_ID: self.unique_id,
-                ATTR_CLUSTER_ID: self.cluster.cluster_id,
-                ATTR_COMMAND: command,
-                ATTR_ARGS: args,
+                EVENT: DeviceEvents.ZHA_EVENT,
+                EVENT_TYPE: EventTypes.DEVICE_EVENT,
+                "command": command,
+                "args": args,
             }
         )
-        """
 
     async def async_update(self) -> None:
         """Retrieve latest state from cluster."""
@@ -512,18 +510,15 @@ class ClientClusterHandler(ClusterHandler):
 
     def attribute_updated(self, attrid: int, value: Any) -> None:
         """Handle an attribute updated on this cluster."""
-        """ TODO
+
         self.zha_send_event(
             SIGNAL_ATTR_UPDATED,
             {
-                ATTR_ATTRIBUTE_ID: attrid,
-                ATTR_ATTRIBUTE_NAME: self._cluster.attributes.get(attrid, ["Unknown"])[
-                    0
-                ],
-                ATTR_VALUE: value,
+                "attribute_id": attrid,
+                "attribute_name": self._cluster.attributes.get(attrid, ["Unknown"])[0],
+                "value": value,
             },
         )
-        """
 
     def cluster_command(self, tsn: int, command_id: int, args: Any) -> None:
         """Handle a cluster command received on this cluster."""
diff --git a/zhaws/server/zigbee/cluster/general.py b/zhaws/server/zigbee/cluster/general.py
index ef68f6a6..c077f6c0 100644
--- a/zhaws/server/zigbee/cluster/general.py
+++ b/zhaws/server/zigbee/cluster/general.py
@@ -424,7 +424,7 @@ def cluster_command(
         """Handle commands received to this cluster."""
         cmd_name = self.cluster.client_commands.get(command_id, [command_id])[0]
         self.debug("Received %s tsn command '%s': %s", tsn, cmd_name, args)
-        # TODO self.zha_send_event(cmd_name, args)
+        self.zha_send_event(cmd_name, args or [])
         if cmd_name == "checkin":
             self.cluster.create_catching_task(self.check_in_response(tsn))
 
diff --git a/zhaws/server/zigbee/cluster/security.py b/zhaws/server/zigbee/cluster/security.py
index 62f13cda..efdc1e31 100644
--- a/zhaws/server/zigbee/cluster/security.py
+++ b/zhaws/server/zigbee/cluster/security.py
@@ -123,7 +123,6 @@ def arm(self, arm_mode: int, code: str | None, zone_id: int) -> None:
         """Handle the IAS ACE arm command."""
         mode = AceCluster.ArmMode(arm_mode)
 
-        """TODO figure out events
         self.zha_send_event(
             self._cluster.server_commands.get(IAS_ACE_ARM)[NAME],
             {
@@ -133,7 +132,6 @@ def arm(self, arm_mode: int, code: str | None, zone_id: int) -> None:
                 "zone_id": zone_id,
             },
         )
-        """
 
         zigbee_reply = self.arm_map[mode](code)
         asyncio.create_task(zigbee_reply)
@@ -218,12 +216,11 @@ def _handle_arm(
 
     def _bypass(self, zone_list: Any, code: str) -> asyncio.Future:
         """Handle the IAS ACE bypass command."""
-        """TODO figure out events
+
         self.zha_send_event(
             self._cluster.server_commands.get(IAS_ACE_BYPASS)[NAME],
             {"zone_list": zone_list, "code": code},
         )
-        """
 
     def _emergency(self) -> asyncio.Future:
         """Handle the IAS ACE emergency command."""
diff --git a/zhaws/server/zigbee/device.py b/zhaws/server/zigbee/device.py
index a9badac8..0c148e23 100644
--- a/zhaws/server/zigbee/device.py
+++ b/zhaws/server/zigbee/device.py
@@ -18,7 +18,6 @@
 from zigpy.zcl.clusters.general import Groups
 import zigpy.zdo.types as zdo_types
 
-from zhaws.backports.enum import StrEnum
 from zhaws.server.const import (
     DEVICE,
     EVENT,
@@ -26,6 +25,7 @@
     IEEE,
     MESSAGE_TYPE,
     NWK,
+    DeviceEvents,
     EventTypes,
     MessageTypes,
 )
@@ -108,13 +108,6 @@ class DeviceStatus(Enum):
     INITIALIZED = 2
 
 
-class DeviceEvents(StrEnum):
-    """Events that devices can broadcast."""
-
-    DEVICE_OFFLINE = "device_offline"
-    DEVICE_ONLINE = "device_online"
-
-
 class Device(LogMixin):
     """ZHAWSS Zigbee device object."""
 
@@ -312,7 +305,10 @@ def device_automation_triggers(self) -> dict:
         if hasattr(self._zigpy_device, "device_automation_triggers"):
             triggers.update(self._zigpy_device.device_automation_triggers)
 
-        return triggers
+        return_triggers = {
+            f"{key[0]}~{key[1]}": value for key, value in triggers.items()
+        }
+        return return_triggers
 
     @property
     def available(self) -> bool:
@@ -574,6 +570,8 @@ def zha_device_info(self) -> dict:
                 )
         device_info[ATTR_ENDPOINT_NAMES] = names
 
+        device_info["device_automation_triggers"] = self.device_automation_triggers
+
         return device_info
 
     def async_get_clusters(self) -> dict[int, dict[CLUSTER_TYPE, list[int]]]:
diff --git a/zhaws/server/zigbee/endpoint.py b/zhaws/server/zigbee/endpoint.py
index 9d82dc7d..69c4c2d5 100644
--- a/zhaws/server/zigbee/endpoint.py
+++ b/zhaws/server/zigbee/endpoint.py
@@ -202,8 +202,8 @@ def send_event(self, signal: dict[str, Any]) -> None:
         """Broadcast an event from this endpoint."""
         signal["endpoint"] = {
             "id": self.id,
+            "unique_id": self.unique_id,
         }
-        # signal["endpoint_id"] = self.id
         self.device.send_event(signal)
 
     def claim_cluster_handlers(self, cluster_handlers: list[ClusterHandler]) -> None:
@@ -218,15 +218,3 @@ def unclaimed_cluster_handlers(self) -> list[ClusterHandler]:
             self.all_cluster_handlers[cluster_id]
             for cluster_id in (available - claimed)
         ]
-
-    """ TODO
-    def zha_send_event(self, event_data: dict[str, str | int]) -> None:
-        #Relay events to hass.
-        self._channels.zha_send_event(
-            {
-                const.ATTR_UNIQUE_ID: self.unique_id,
-                const.ATTR_ENDPOINT_ID: self.id,
-                **event_data,
-            }
-        )
-    """