From 987c7a289ab678a71f7c887383435217642832e1 Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Tue, 13 Jul 2021 13:31:17 -0400
Subject: [PATCH] Update ZHA to support zigpy 0.34.0 device initialization
 (#52610)

* Handle `None` node descriptors

* Skip loading uninitialized devices

* Fix unit test incorrectly handling unset cluster `ep_attribute`

* Revert filtering devices by status during startup
---
 homeassistant/components/zha/core/device.py | 39 ++++++++++++++-------
 tests/components/zha/common.py              | 17 +++++----
 2 files changed, 34 insertions(+), 22 deletions(-)

diff --git a/homeassistant/components/zha/core/device.py b/homeassistant/components/zha/core/device.py
index 0c572bfba8a..c6166419e39 100644
--- a/homeassistant/components/zha/core/device.py
+++ b/homeassistant/components/zha/core/device.py
@@ -183,11 +183,12 @@ class ZHADevice(LogMixin):
         return self._zigpy_device.model
 
     @property
-    def manufacturer_code(self):
+    def manufacturer_code(self) -> int | None:
         """Return the manufacturer code for the device."""
-        if self._zigpy_device.node_desc.is_valid:
-            return self._zigpy_device.node_desc.manufacturer_code
-        return None
+        if self._zigpy_device.node_desc is None:
+            return None
+
+        return self._zigpy_device.node_desc.manufacturer_code
 
     @property
     def nwk(self):
@@ -210,17 +211,20 @@ class ZHADevice(LogMixin):
         return self._zigpy_device.last_seen
 
     @property
-    def is_mains_powered(self):
+    def is_mains_powered(self) -> bool | None:
         """Return true if device is mains powered."""
+        if self._zigpy_device.node_desc is None:
+            return None
+
         return self._zigpy_device.node_desc.is_mains_powered
 
     @property
-    def device_type(self):
+    def device_type(self) -> str:
         """Return the logical device type for the device."""
-        node_descriptor = self._zigpy_device.node_desc
-        return (
-            node_descriptor.logical_type.name if node_descriptor.is_valid else UNKNOWN
-        )
+        if self._zigpy_device.node_desc is None:
+            return UNKNOWN
+
+        return self._zigpy_device.node_desc.logical_type.name
 
     @property
     def power_source(self):
@@ -230,18 +234,27 @@ class ZHADevice(LogMixin):
         )
 
     @property
-    def is_router(self):
+    def is_router(self) -> bool | None:
         """Return true if this is a routing capable device."""
+        if self._zigpy_device.node_desc is None:
+            return None
+
         return self._zigpy_device.node_desc.is_router
 
     @property
-    def is_coordinator(self):
+    def is_coordinator(self) -> bool | None:
         """Return true if this device represents the coordinator."""
+        if self._zigpy_device.node_desc is None:
+            return None
+
         return self._zigpy_device.node_desc.is_coordinator
 
     @property
-    def is_end_device(self):
+    def is_end_device(self) -> bool | None:
         """Return true if this device is an end device."""
+        if self._zigpy_device.node_desc is None:
+            return None
+
         return self._zigpy_device.node_desc.is_end_device
 
     @property
diff --git a/tests/components/zha/common.py b/tests/components/zha/common.py
index eb65cc4fd2e..5180e9dbc07 100644
--- a/tests/components/zha/common.py
+++ b/tests/components/zha/common.py
@@ -3,8 +3,8 @@ import asyncio
 import time
 from unittest.mock import AsyncMock, Mock
 
-from zigpy.device import Device as zigpy_dev
-from zigpy.endpoint import Endpoint as zigpy_ep
+import zigpy.device as zigpy_dev
+import zigpy.endpoint as zigpy_ep
 import zigpy.profiles.zha
 import zigpy.types
 import zigpy.zcl
@@ -27,7 +27,7 @@ class FakeEndpoint:
         self.out_clusters = {}
         self._cluster_attr = {}
         self.member_of = {}
-        self.status = 1
+        self.status = zigpy_ep.Status.ZDO_INIT
         self.manufacturer = manufacturer
         self.model = model
         self.profile_id = zigpy.profiles.zha.PROFILE_ID
@@ -57,7 +57,7 @@ class FakeEndpoint:
     @property
     def __class__(self):
         """Fake being Zigpy endpoint."""
-        return zigpy_ep
+        return zigpy_ep.Endpoint
 
     @property
     def unique_id(self):
@@ -65,8 +65,8 @@ class FakeEndpoint:
         return self.device.ieee, self.endpoint_id
 
 
-FakeEndpoint.add_to_group = zigpy_ep.add_to_group
-FakeEndpoint.remove_from_group = zigpy_ep.remove_from_group
+FakeEndpoint.add_to_group = zigpy_ep.Endpoint.add_to_group
+FakeEndpoint.remove_from_group = zigpy_ep.Endpoint.remove_from_group
 
 
 def patch_cluster(cluster):
@@ -125,12 +125,11 @@ class FakeDevice:
         self.lqi = 255
         self.rssi = 8
         self.last_seen = time.time()
-        self.status = 2
+        self.status = zigpy_dev.Status.ENDPOINTS_INIT
         self.initializing = False
         self.skip_configuration = False
         self.manufacturer = manufacturer
         self.model = model
-        self.node_desc = zigpy.zdo.types.NodeDescriptor()
         self.remove_from_group = AsyncMock()
         if node_desc is None:
             node_desc = b"\x02@\x807\x10\x7fd\x00\x00*d\x00\x00"
@@ -138,7 +137,7 @@ class FakeDevice:
         self.neighbors = []
 
 
-FakeDevice.add_to_group = zigpy_dev.add_to_group
+FakeDevice.add_to_group = zigpy_dev.Device.add_to_group
 
 
 def get_zha_gateway(hass):