Reorganize ZHA device availability code (#108856)

* Correct ZHA device availability at startup

* don't set available property from gateway

* cleanup

homeassistant/components/zha/core/device.py

@@ -11,7 +11,7 @@ import time
 from typing import TYPE_CHECKING, Any, Self
 
 from zigpy import types
-import zigpy.device
+from zigpy.device import Device as ZigpyDevice
 import zigpy.exceptions
 from zigpy.profiles import PROFILES
 import zigpy.quirks
@@ -124,22 +124,23 @@ class ZHADevice(LogMixin):
         zha_gateway: ZHAGateway,
     ) -> None:
         """Initialize the gateway."""
-        self.hass = hass
-        self._zigpy_device = zigpy_device
-        self._zha_gateway = zha_gateway
-        self._available = False
-        self._available_signal = f"{self.name}_{self.ieee}_{SIGNAL_AVAILABLE}"
-        self._checkins_missed_count = 0
+        self.hass: HomeAssistant = hass
+        self._zigpy_device: ZigpyDevice = zigpy_device
+        self._zha_gateway: ZHAGateway = zha_gateway
+        self._available_signal: str = f"{self.name}_{self.ieee}_{SIGNAL_AVAILABLE}"
+        self._checkins_missed_count: int = 0
         self.unsubs: list[Callable[[], None]] = []
-        self.quirk_applied = isinstance(self._zigpy_device, zigpy.quirks.CustomDevice)
-        self.quirk_class = (
+        self.quirk_applied: bool = isinstance(
+            self._zigpy_device, zigpy.quirks.CustomDevice
+        )
+        self.quirk_class: str = (
             f"{self._zigpy_device.__class__.__module__}."
             f"{self._zigpy_device.__class__.__name__}"
         )
-        self.quirk_id = getattr(self._zigpy_device, ATTR_QUIRK_ID, None)
+        self.quirk_id: str | None = getattr(self._zigpy_device, ATTR_QUIRK_ID, None)
 
         if self.is_mains_powered:
-            self.consider_unavailable_time = async_get_zha_config_value(
+            self.consider_unavailable_time: int = async_get_zha_config_value(
                 self._zha_gateway.config_entry,
                 ZHA_OPTIONS,
                 CONF_CONSIDER_UNAVAILABLE_MAINS,
@@ -152,7 +153,10 @@ class ZHADevice(LogMixin):
                 CONF_CONSIDER_UNAVAILABLE_BATTERY,
                 CONF_DEFAULT_CONSIDER_UNAVAILABLE_BATTERY,
             )
-
+        self._available: bool = self.is_coordinator or (
+            self.last_seen is not None
+            and time.time() - self.last_seen < self.consider_unavailable_time
+        )
         self._zdo_handler: ZDOClusterHandler = ZDOClusterHandler(self)
         self._power_config_ch: ClusterHandler | None = None
         self._identify_ch: ClusterHandler | None = None
@@ -408,7 +412,6 @@ class ZHADevice(LogMixin):
         hass: HomeAssistant,
         zigpy_dev: zigpy.device.Device,
         gateway: ZHAGateway,
-        restored: bool = False,
     ) -> Self:
         """Create new device."""
         zha_dev = cls(hass, zigpy_dev, gateway)

homeassistant/components/zha/core/gateway.py

@@ -223,7 +223,7 @@ class ZHAGateway:
         zha_data.gateway = self
 
         self.coordinator_zha_device = self._async_get_or_create_device(
-            self._find_coordinator_device(), restored=True
+            self._find_coordinator_device()
         )
 
         self.async_load_devices()
@@ -264,11 +264,10 @@ class ZHAGateway:
         """Restore ZHA devices from zigpy application state."""
 
         for zigpy_device in self.application_controller.devices.values():
-            zha_device = self._async_get_or_create_device(zigpy_device, restored=True)
+            zha_device = self._async_get_or_create_device(zigpy_device)
             delta_msg = "not known"
             if zha_device.last_seen is not None:
                 delta = round(time.time() - zha_device.last_seen)
-                zha_device.available = delta < zha_device.consider_unavailable_time
                 delta_msg = f"{str(timedelta(seconds=delta))} ago"
             _LOGGER.debug(
                 (
@@ -622,11 +621,11 @@ class ZHAGateway:
 
     @callback
     def _async_get_or_create_device(
-        self, zigpy_device: zigpy.device.Device, restored: bool = False
+        self, zigpy_device: zigpy.device.Device
     ) -> ZHADevice:
         """Get or create a ZHA device."""
         if (zha_device := self._devices.get(zigpy_device.ieee)) is None:
-            zha_device = ZHADevice.new(self.hass, zigpy_device, self, restored)
+            zha_device = ZHADevice.new(self.hass, zigpy_device, self)
             self._devices[zigpy_device.ieee] = zha_device
 
             device_registry = dr.async_get(self.hass)

tests/components/zha/conftest.py

@@ -25,6 +25,7 @@ import zigpy.zdo.types as zdo_t
 
 import homeassistant.components.zha.core.const as zha_const
 import homeassistant.components.zha.core.device as zha_core_device
+from homeassistant.components.zha.core.gateway import ZHAGateway
 from homeassistant.components.zha.core.helpers import get_zha_gateway
 from homeassistant.helpers import restore_state
 from homeassistant.setup import async_setup_component
@@ -381,7 +382,7 @@ def zha_device_joined_restored(request):
 
 @pytest.fixture
 def zha_device_mock(
-    hass, zigpy_device_mock
+    hass, config_entry, zigpy_device_mock
 ) -> Callable[..., zha_core_device.ZHADevice]:
     """Return a ZHA Device factory."""
 
@@ -409,7 +410,11 @@ def zha_device_mock(
         zigpy_device = zigpy_device_mock(
             endpoints, ieee, manufacturer, model, node_desc, patch_cluster=patch_cluster
         )
-        zha_device = zha_core_device.ZHADevice(hass, zigpy_device, MagicMock())
+        zha_device = zha_core_device.ZHADevice(
+            hass,
+            zigpy_device,
+            ZHAGateway(hass, {}, config_entry),
+        )
         return zha_device
 
     return _zha_device

tests/components/zha/test_gateway.py

@@ -365,7 +365,7 @@ async def test_startup_concurrency_limit(
             zigpy.zdo.types.NodeDescriptor.MACCapabilityFlags.MainsPowered
         )
 
-        zha_gateway._async_get_or_create_device(zigpy_dev, restored=True)
+        zha_gateway._async_get_or_create_device(zigpy_dev)
 
     # Keep track of request concurrency during initialization
     current_concurrency = 0