Split ZHA device loading and entities adding (#33075)

* Split ZHA device loading and entities adding.
Load zha devices from ZHA gateway, but add entities after integration
was setup.

* Use hass.loop.create_task()
* Split ZHA device initialization from loading.
Restore and initialize ZHA devices separately.

* Use hass.async_create_task()
Load devices prior group initialization.

homeassistant/components/zha/__init__.py

@@ -8,6 +8,8 @@ import voluptuous as vol
 from homeassistant import config_entries, const as ha_const
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.device_registry import CONNECTION_ZIGBEE
+from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.helpers.typing import HomeAssistantType
 
 from . import api
 from .core import ZHAGateway
@@ -27,6 +29,7 @@ from .core.const import (
     DEFAULT_BAUDRATE,
     DEFAULT_RADIO_TYPE,
     DOMAIN,
+    SIGNAL_ADD_ENTITIES,
     RadioType,
 )
 
@@ -90,8 +93,15 @@ async def async_setup_entry(hass, config_entry):
     """
 
     zha_data = hass.data.setdefault(DATA_ZHA, {})
+    zha_data[DATA_ZHA_PLATFORM_LOADED] = {}
     config = zha_data.get(DATA_ZHA_CONFIG, {})
 
+    zha_data[DATA_ZHA_DISPATCHERS] = []
+    for component in COMPONENTS:
+        zha_data[component] = []
+        coro = hass.config_entries.async_forward_entry_setup(config_entry, component)
+        zha_data[DATA_ZHA_PLATFORM_LOADED][component] = hass.async_create_task(coro)
+
     if config.get(CONF_ENABLE_QUIRKS, True):
         # needs to be done here so that the ZHA module is finished loading
         # before zhaquirks is imported
@@ -100,22 +110,6 @@ async def async_setup_entry(hass, config_entry):
     zha_gateway = ZHAGateway(hass, config, config_entry)
     await zha_gateway.async_initialize()
 
-    zha_data[DATA_ZHA_DISPATCHERS] = []
-    zha_data[DATA_ZHA_PLATFORM_LOADED] = asyncio.Event()
-    platforms = []
-    for component in COMPONENTS:
-        platforms.append(
-            hass.async_create_task(
-                hass.config_entries.async_forward_entry_setup(config_entry, component)
-            )
-        )
-
-    async def _platforms_loaded():
-        await asyncio.gather(*platforms)
-        zha_data[DATA_ZHA_PLATFORM_LOADED].set()
-
-    hass.async_create_task(_platforms_loaded())
-
     device_registry = await hass.helpers.device_registry.async_get_registry()
     device_registry.async_get_or_create(
         config_entry_id=config_entry.entry_id,
@@ -134,7 +128,7 @@ async def async_setup_entry(hass, config_entry):
         await zha_data[DATA_ZHA_GATEWAY].async_update_device_storage()
 
     hass.bus.async_listen_once(ha_const.EVENT_HOMEASSISTANT_STOP, async_zha_shutdown)
-    hass.async_create_task(zha_gateway.async_load_devices())
+    hass.async_create_task(async_load_entities(hass, config_entry))
     return True
 
 
@@ -152,3 +146,20 @@ async def async_unload_entry(hass, config_entry):
         await hass.config_entries.async_forward_entry_unload(config_entry, component)
 
     return True
+
+
+async def async_load_entities(
+    hass: HomeAssistantType, config_entry: config_entries.ConfigEntry
+) -> None:
+    """Load entities after integration was setup."""
+    await hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].async_prepare_entities()
+    to_setup = [
+        hass.data[DATA_ZHA][DATA_ZHA_PLATFORM_LOADED][comp]
+        for comp in COMPONENTS
+        if hass.data[DATA_ZHA][comp]
+    ]
+    results = await asyncio.gather(*to_setup, return_exceptions=True)
+    for res in results:
+        if isinstance(res, Exception):
+            _LOGGER.warning("Couldn't setup zha platform: %s", res)
+    async_dispatcher_send(hass, SIGNAL_ADD_ENTITIES)

homeassistant/components/zha/binary_sensor.py

@@ -49,7 +49,7 @@ STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, DOMAIN)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation binary sensor from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/core/discovery.py

@@ -8,6 +8,16 @@ from homeassistant.core import callback
 from homeassistant.helpers.typing import HomeAssistantType
 
 from . import const as zha_const, registries as zha_regs, typing as zha_typing
+from .. import (  # noqa: F401 pylint: disable=unused-import,
+    binary_sensor,
+    cover,
+    device_tracker,
+    fan,
+    light,
+    lock,
+    sensor,
+    switch,
+)
 from .channels import base
 
 _LOGGER = logging.getLogger(__name__)

homeassistant/components/zha/core/gateway.py

@@ -36,7 +36,6 @@ from .const import (
     DATA_ZHA,
     DATA_ZHA_BRIDGE_ID,
     DATA_ZHA_GATEWAY,
-    DATA_ZHA_PLATFORM_LOADED,
     DEBUG_COMP_BELLOWS,
     DEBUG_COMP_ZHA,
     DEBUG_COMP_ZIGPY,
@@ -157,34 +156,40 @@ class ZHAGateway:
         self._hass.data[DATA_ZHA][DATA_ZHA_BRIDGE_ID] = str(
             self.application_controller.ieee
         )
+        await self.async_load_devices()
         self._initialize_groups()
 
     async def async_load_devices(self) -> None:
         """Restore ZHA devices from zigpy application state."""
-        await self._hass.data[DATA_ZHA][DATA_ZHA_PLATFORM_LOADED].wait()
+        zigpy_devices = self.application_controller.devices.values()
+        for zigpy_device in zigpy_devices:
+            self._async_get_or_create_device(zigpy_device, restored=True)
 
+    async def async_prepare_entities(self) -> None:
+        """Prepare entities by initializing device channels."""
         semaphore = asyncio.Semaphore(2)
 
-        async def _throttle(device: zha_typing.ZigpyDeviceType):
+        async def _throttle(zha_device: zha_typing.ZhaDeviceType, cached: bool):
             async with semaphore:
-                await self.async_device_restored(device)
+                await zha_device.async_initialize(from_cache=cached)
 
-        zigpy_devices = self.application_controller.devices.values()
         _LOGGER.debug("Loading battery powered devices")
         await asyncio.gather(
             *[
-                _throttle(dev)
+                _throttle(dev, cached=True)
-                for dev in zigpy_devices
+                for dev in self.devices.values()
-                if not dev.node_desc.is_mains_powered
+                if not dev.is_mains_powered
             ]
         )
-        async_dispatcher_send(self._hass, SIGNAL_ADD_ENTITIES)
 
         _LOGGER.debug("Loading mains powered devices")
         await asyncio.gather(
-            *[_throttle(dev) for dev in zigpy_devices if dev.node_desc.is_mains_powered]
+            *[
+                _throttle(dev, cached=False)
+                for dev in self.devices.values()
+                if dev.is_mains_powered
+            ]
         )
-        async_dispatcher_send(self._hass, SIGNAL_ADD_ENTITIES)
 
     def device_joined(self, device):
         """Handle device joined.

homeassistant/components/zha/cover.py

@@ -29,7 +29,7 @@ STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, DOMAIN)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation cover from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/device_tracker.py

@@ -26,7 +26,7 @@ _LOGGER = logging.getLogger(__name__)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation device tracker from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/fan.py

@@ -53,7 +53,7 @@ STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, DOMAIN)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation fan from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/light.py

@@ -51,7 +51,7 @@ PARALLEL_UPDATES = 0
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation light from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][light.DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][light.DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/lock.py

@@ -36,7 +36,7 @@ VALUE_TO_STATE = dict(enumerate(STATE_LIST))
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation Door Lock from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/sensor.py

@@ -68,7 +68,7 @@ STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, DOMAIN)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation sensor from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,

homeassistant/components/zha/switch.py

@@ -26,7 +26,7 @@ STRICT_MATCH = functools.partial(ZHA_ENTITIES.strict_match, DOMAIN)
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up the Zigbee Home Automation switch from config entry."""
-    entities_to_create = hass.data[DATA_ZHA][DOMAIN] = []
+    entities_to_create = hass.data[DATA_ZHA][DOMAIN]
 
     unsub = async_dispatcher_connect(
         hass,