use enums

Author: dmulcahey

zha/application/model.py

@@ -5,7 +5,9 @@
 
 from zigpy.types.named import EUI64, NWK
 
+from zha.const import EventTypes
 from zha.model import BaseEvent, BaseModel
+from zha.websocket.const import ControllerEvents, DeviceEvents
 from zha.zigbee.model import DeviceInfo, ExtendedDeviceInfo, GroupInfo
 
 
@@ -41,26 +43,26 @@ class DeviceJoinedDeviceInfo(BaseModel):
 class ConnectionLostEvent(BaseEvent):
     """Event to signal that the connection to the radio has been lost."""
 
-    event_type: Literal["zha_gateway_message"] = "zha_gateway_message"
-    event: Literal["connection_lost"] = "connection_lost"
+    event_type: Literal[EventTypes.CONTROLLER_EVENT] = EventTypes.CONTROLLER_EVENT
+    event: Literal[ControllerEvents.CONNECTION_LOST] = ControllerEvents.CONNECTION_LOST
     exception: Exception | None = None
 
 
 class DeviceJoinedEvent(BaseEvent):
     """Event to signal that a device has joined the network."""
 
     device_info: DeviceJoinedDeviceInfo
-    event_type: Literal["zha_gateway_message"] = "zha_gateway_message"
-    event: Literal["device_joined"] = "device_joined"
+    event_type: Literal[EventTypes.CONTROLLER_EVENT] = EventTypes.CONTROLLER_EVENT
+    event: Literal[ControllerEvents.DEVICE_JOINED] = ControllerEvents.DEVICE_JOINED
 
 
 class DeviceLeftEvent(BaseEvent):
     """Event to signal that a device has left the network."""
 
     ieee: EUI64
     nwk: NWK
-    event_type: Literal["zha_gateway_message"] = "zha_gateway_message"
-    event: Literal["device_left"] = "device_left"
+    event_type: Literal[EventTypes.CONTROLLER_EVENT] = EventTypes.CONTROLLER_EVENT
+    event: Literal[ControllerEvents.DEVICE_LEFT] = ControllerEvents.DEVICE_LEFT
 
 
 class RawDeviceInitializedDeviceInfo(DeviceJoinedDeviceInfo):
@@ -75,70 +77,78 @@ class RawDeviceInitializedEvent(BaseEvent):
     """Event to signal that a device has been initialized without quirks loaded."""
 
     device_info: RawDeviceInitializedDeviceInfo
-    event_type: Literal["zha_gateway_message"] = "zha_gateway_message"
-    event: Literal["raw_device_initialized"] = "raw_device_initialized"
+    event_type: Literal[EventTypes.CONTROLLER_EVENT] = EventTypes.CONTROLLER_EVENT
+    event: Literal[ControllerEvents.RAW_DEVICE_INITIALIZED] = (
+        ControllerEvents.RAW_DEVICE_INITIALIZED
+    )
 
 
 class DeviceFullyInitializedEvent(BaseEvent):
     """Event to signal that a device has been fully initialized."""
 
     device_info: ExtendedDeviceInfoWithPairingStatus
     new_join: bool = False
-    event_type: Literal["zha_gateway_message"] = "zha_gateway_message"
-    event: Literal["device_fully_initialized"] = "device_fully_initialized"
+    event_type: Literal[EventTypes.CONTROLLER_EVENT] = EventTypes.CONTROLLER_EVENT
+    event: Literal[ControllerEvents.DEVICE_FULLY_INITIALIZED] = (
+        ControllerEvents.DEVICE_FULLY_INITIALIZED
+    )
 
 
 class GroupRemovedEvent(BaseEvent):
     """Group removed event."""
 
-    event_type: Literal["zha_gateway_message"] = "zha_gateway_message"
-    event: Literal["group_removed"] = "group_removed"
+    event_type: Literal[EventTypes.CONTROLLER_EVENT] = EventTypes.CONTROLLER_EVENT
+    event: Literal[ControllerEvents.GROUP_REMOVED] = ControllerEvents.GROUP_REMOVED
     group_info: GroupInfo
 
 
 class GroupAddedEvent(BaseEvent):
     """Group added event."""
 
-    event_type: Literal["zha_gateway_message"] = "zha_gateway_message"
-    event: Literal["group_added"] = "group_added"
+    event_type: Literal[EventTypes.CONTROLLER_EVENT] = EventTypes.CONTROLLER_EVENT
+    event: Literal[ControllerEvents.GROUP_ADDED] = ControllerEvents.GROUP_ADDED
     group_info: GroupInfo
 
 
 class GroupMemberAddedEvent(BaseEvent):
     """Group member added event."""
 
-    event_type: Literal["zha_gateway_message"] = "zha_gateway_message"
-    event: Literal["group_member_added"] = "group_member_added"
+    event_type: Literal[EventTypes.CONTROLLER_EVENT] = EventTypes.CONTROLLER_EVENT
+    event: Literal[ControllerEvents.GROUP_MEMBER_ADDED] = (
+        ControllerEvents.GROUP_MEMBER_ADDED
+    )
     group_info: GroupInfo
 
 
 class GroupMemberRemovedEvent(BaseEvent):
     """Group member removed event."""
 
-    event_type: Literal["zha_gateway_message"] = "zha_gateway_message"
-    event: Literal["group_member_removed"] = "group_member_removed"
+    event_type: Literal[EventTypes.CONTROLLER_EVENT] = EventTypes.CONTROLLER_EVENT
+    event: Literal[ControllerEvents.GROUP_MEMBER_REMOVED] = (
+        ControllerEvents.GROUP_MEMBER_REMOVED
+    )
     group_info: GroupInfo
 
 
 class DeviceRemovedEvent(BaseEvent):
     """Event to signal that a device has been removed."""
 
     device_info: ExtendedDeviceInfo
-    event_type: Literal["zha_gateway_message"] = "zha_gateway_message"
-    event: Literal["device_removed"] = "device_removed"
+    event_type: Literal[EventTypes.CONTROLLER_EVENT] = EventTypes.CONTROLLER_EVENT
+    event: Literal[ControllerEvents.DEVICE_REMOVED] = ControllerEvents.DEVICE_REMOVED
 
 
 class DeviceOfflineEvent(BaseEvent):
     """Device offline event."""
 
-    event: Literal["device_offline"] = "device_offline"
-    event_type: Literal["device_event"] = "device_event"
+    event: Literal[DeviceEvents.DEVICE_OFFLINE] = DeviceEvents.DEVICE_OFFLINE
+    event_type: Literal[EventTypes.DEVICE_EVENT] = EventTypes.DEVICE_EVENT
     device_info: ExtendedDeviceInfo
 
 
 class DeviceOnlineEvent(BaseEvent):
     """Device online event."""
 
-    event: Literal["device_online"] = "device_online"
-    event_type: Literal["device_event"] = "device_event"
+    event: Literal[DeviceEvents.DEVICE_ONLINE] = DeviceEvents.DEVICE_ONLINE
+    event_type: Literal[EventTypes.DEVICE_EVENT] = EventTypes.DEVICE_EVENT
     device_info: ExtendedDeviceInfo

zha/websocket/const.py

@@ -104,15 +104,6 @@ class MessageTypes(StrEnum):
     RESULT = "result"
 
 
-class EventTypes(StrEnum):
-    """WS event types."""
-
-    CONTROLLER_EVENT = "zha_gateway_message"
-    PLATFORM_ENTITY_EVENT = "platform_entity_event"
-    RAW_ZCL_EVENT = "raw_zcl_event"
-    DEVICE_EVENT = "device_event"
-
-
 class ControllerEvents(StrEnum):
     """WS controller events."""
 
@@ -126,6 +117,7 @@ class ControllerEvents(StrEnum):
     GROUP_MEMBER_REMOVED = "group_member_removed"
     GROUP_ADDED = "group_added"
     GROUP_REMOVED = "group_removed"
+    CONNECTION_LOST = "connection_lost"
 
 
 class PlatformEntityEvents(StrEnum):

zha/websocket/server/client.py

@@ -10,6 +10,7 @@
 from pydantic import BaseModel, ValidationError
 from websockets.server import WebSocketServerProtocol
 
+from zha.const import EventTypes
 from zha.model import BaseEvent
 from zha.websocket.const import (
     COMMAND,
@@ -22,7 +23,6 @@
     ZIGBEE_ERROR,
     ZIGBEE_ERROR_CODE,
     APICommands,
-    EventTypes,
     MessageTypes,
 )
 from zha.websocket.server.api import decorators, register_api_command