Update entity state when ZHA device becomes available (#20993)

* correctly update device entity state

* update state when device becomes available

* constants

* review comments

homeassistant/components/zha/core/device.py

@@ -173,7 +173,7 @@ class ZHADevice:
         await self._execute_listener_tasks('async_configure')
         _LOGGER.debug('%s: completed configuration', self.name)
 
-    async def async_initialize(self, from_cache):
+    async def async_initialize(self, from_cache=False):
         """Initialize listeners."""
         _LOGGER.debug('%s: started initialization', self.name)
         await self._execute_listener_tasks('async_initialize', from_cache)

homeassistant/components/zha/core/gateway.py

@@ -138,7 +138,7 @@ class ZHAGateway:
             ))
         await asyncio.gather(*endpoint_tasks)
 
-        await zha_device.async_initialize(not is_new_join)
+        await zha_device.async_initialize(from_cache=(not is_new_join))
 
         discovery_tasks = []
         for discovery_info in discovery_infos:

homeassistant/components/zha/device_entity.py

@@ -8,6 +8,7 @@ https://home-assistant.io/components/zha/
 import logging
 import time
 
+from homeassistant.core import callback
 from homeassistant.util import slugify
 from .entity import ZhaEntity
 from .const import LISTENER_BATTERY, SIGNAL_STATE_ATTR
@@ -30,6 +31,9 @@ BATTERY_SIZES = {
     255: 'Unknown'
 }
 
+STATE_ONLINE = 'online'
+STATE_OFFLINE = 'offline'
+
 
 class ZhaDeviceEntity(ZhaEntity):
     """A base class for ZHA devices."""
@@ -108,13 +112,20 @@ class ZhaDeviceEntity(ZhaEntity):
             difference = time.time() - self._zha_device.last_seen
             if difference > self._keepalive_interval:
                 self._zha_device.update_available(False)
-                self._state = None
             else:
                 self._zha_device.update_available(True)
-                self._state = 'online'
                 if self._battery_listener:
                     await self.async_get_latest_battery_reading()
 
+    @callback
+    def async_set_available(self, available):
+        """Set entity availability."""
+        if available:
+            self._state = STATE_ONLINE
+        else:
+            self._state = STATE_OFFLINE
+        super().async_set_available(available)
+
     async def _async_init_battery_values(self):
         """Get initial battery level and battery info from listener cache."""
         battery_size = await self._battery_listener.get_attribute_value(