From e0f99385b9341f9a658d53b4265d67785e5ebe58 Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Thu, 8 Feb 2024 23:06:56 -0500
Subject: [PATCH] Fix startup

---
 bellows/ezsp/__init__.py      |  4 ++--
 bellows/zigbee/application.py | 10 +++++-----
 bellows/zigbee/device.py      | 26 +++++++++++++++-----------
 3 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/bellows/ezsp/__init__.py b/bellows/ezsp/__init__.py
index 6e88755d..f2f8d6cf 100644
--- a/bellows/ezsp/__init__.py
+++ b/bellows/ezsp/__init__.py
@@ -644,7 +644,7 @@ async def get_supported_firmware_features(
                 bytes([custom_commands.CustomCommand.CMD_GET_SUPPORTED_FEATURES]),
             )
         except InvalidCommandError:
-            return custom_commands.SupportedCustomFeatures(0)
+            return custom_commands.FirmwareFeatures(0)
 
-        features, _ = custom_commands.SupportedCustomFeatures.deserialize(rsp_data)
+        features, _ = custom_commands.FirmwareFeatures.deserialize(rsp_data)
         return features
diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py
index 64921860..892b2470 100644
--- a/bellows/zigbee/application.py
+++ b/bellows/zigbee/application.py
@@ -37,7 +37,7 @@
 import bellows.multicast
 import bellows.types as t
 from bellows.zigbee import repairs
-from bellows.zigbee.device import EZSPEndpoint
+from bellows.zigbee.device import EZSPEndpoint, EZSPGroupEndpoint
 import bellows.zigbee.util as util
 
 APS_ACK_TIMEOUT = 120
@@ -205,14 +205,14 @@ async def start_network(self):
         ezsp.add_callback(self.ezsp_callback_handler)
         self.controller_event.set()
 
-        custom_features = await self._ezsp.get_supported_custom_features()
+        custom_features = await self._ezsp.get_supported_firmware_features()
         LOGGER.debug("Supported custom firmware features: %r", custom_features)
 
         if FirmwareFeatures.MEMBER_OF_ALL_GROUPS in custom_features:
             # If the firmware passes through all incoming group messages, do nothing
-            endpoint_cls = zigpy.endpoint.Endpoint
-        else:
             endpoint_cls = EZSPEndpoint
+        else:
+            endpoint_cls = EZSPGroupEndpoint
 
         ezsp_device = zigpy.device.Device(
             application=self,
@@ -224,7 +224,7 @@ async def start_network(self):
         # The coordinator device does not respond to attribute reads so we have to
         # divine the internal NCP state.
         for zdo_desc in self._created_device_endpoints:
-            ep = endpoint_cls(ezsp_device, zdo_desc.endpoint, zdo_desc)
+            ep = endpoint_cls.from_descriptor(ezsp_device, zdo_desc.endpoint, zdo_desc)
             ezsp_device.endpoints[zdo_desc.endpoint] = ep
             ezsp_device.model = ep.model
             ezsp_device.manufacturer = ep.manufacturer
diff --git a/bellows/zigbee/device.py b/bellows/zigbee/device.py
index 9e4a6e7b..6de7f2dd 100644
--- a/bellows/zigbee/device.py
+++ b/bellows/zigbee/device.py
@@ -21,30 +21,32 @@
 
 
 class EZSPEndpoint(zigpy.endpoint.Endpoint):
-    def __init__(
-        self,
+    @classmethod
+    def from_descriptor(
+        cls,
         device: zigpy.device.Device,
         endpoint_id: int,
         descriptor: zdo_t.SimpleDescriptor,
     ) -> None:
-        super().__init__(device, endpoint_id)
+        ep = cls(device, endpoint_id)
+        ep.profile_id = descriptor.profile
 
-        self.profile_id = descriptor.profile
-
-        if self.profile_id in PROFILE_TO_DEVICE_TYPE:
-            self.device_type = PROFILE_TO_DEVICE_TYPE[self.profile_id](
+        if ep.profile_id in PROFILE_TO_DEVICE_TYPE:
+            ep.device_type = PROFILE_TO_DEVICE_TYPE[ep.profile_id](
                 descriptor.device_type
             )
         else:
-            self.device_type = descriptor.device_type
+            ep.device_type = descriptor.device_type
 
         for cluster in descriptor.input_clusters:
-            self.add_input_cluster(cluster)
+            ep.add_input_cluster(cluster)
 
         for cluster in descriptor.output_clusters:
-            self.add_output_cluster(cluster)
+            ep.add_output_cluster(cluster)
+
+        ep.status = zigpy.endpoint.Status.ZDO_INIT
 
-        self.status = zigpy.endpoint.Status.ZDO_INIT
+        return ep
 
     @property
     def manufacturer(self) -> str:
@@ -56,6 +58,8 @@ def model(self) -> str:
         """Model."""
         return "EZSP"
 
+
+class EZSPGroupEndpoint(EZSPEndpoint):
     async def add_to_group(self, grp_id: int, name: str = None) -> t.EmberStatus:
         if grp_id in self.member_of:
             return t.EmberStatus.SUCCESS