Refactor ZHA setup (#32959)

* Refactor ZHA setup.
Catch errors and raise if needed.

* Cleanup.

homeassistant/components/zha/__init__.py

@@ -89,9 +89,19 @@ async def async_setup_entry(hass, config_entry):
     Will automatically load components to support devices found on the network.
     """
 
-    hass.data[DATA_ZHA] = hass.data.get(DATA_ZHA, {})
+    zha_data = hass.data.setdefault(DATA_ZHA, {})
-    hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS] = []
+    config = zha_data.get(DATA_ZHA_CONFIG, {})
-    hass.data[DATA_ZHA][DATA_ZHA_PLATFORM_LOADED] = asyncio.Event()
+
+    if config.get(CONF_ENABLE_QUIRKS, True):
+        # needs to be done here so that the ZHA module is finished loading
+        # before zhaquirks is imported
+        import zhaquirks  # noqa: F401 pylint: disable=unused-import, import-outside-toplevel, import-error
+
+    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(
@@ -102,20 +112,10 @@ async def async_setup_entry(hass, config_entry):
 
     async def _platforms_loaded():
         await asyncio.gather(*platforms)
-        hass.data[DATA_ZHA][DATA_ZHA_PLATFORM_LOADED].set()
+        zha_data[DATA_ZHA_PLATFORM_LOADED].set()
 
     hass.async_create_task(_platforms_loaded())
 
-    config = hass.data[DATA_ZHA].get(DATA_ZHA_CONFIG, {})
-
-    if config.get(CONF_ENABLE_QUIRKS, True):
-        # needs to be done here so that the ZHA module is finished loading
-        # before zhaquirks is imported
-        import zhaquirks  # noqa: F401 pylint: disable=unused-import, import-outside-toplevel, import-error
-
-    zha_gateway = ZHAGateway(hass, config, config_entry)
-    await zha_gateway.async_initialize()
-
     device_registry = await hass.helpers.device_registry.async_get_registry()
     device_registry.async_get_or_create(
         config_entry_id=config_entry.entry_id,
@@ -130,8 +130,8 @@ async def async_setup_entry(hass, config_entry):
 
     async def async_zha_shutdown(event):
         """Handle shutdown tasks."""
-        await hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].shutdown()
+        await zha_data[DATA_ZHA_GATEWAY].shutdown()
-        await hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].async_update_device_storage()
+        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())

homeassistant/components/zha/core/gateway.py

@@ -7,10 +7,12 @@ import logging
 import os
 import traceback
 
+from serial import SerialException
 import zigpy.device as zigpy_dev
 
 from homeassistant.components.system_log import LogEntry, _figure_out_source
 from homeassistant.core import callback
+from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.device_registry import (
     CONNECTION_ZIGBEE,
     async_get_registry as get_dev_reg,
@@ -98,7 +100,6 @@ class ZHAGateway:
         self.ha_entity_registry = None
         self.application_controller = None
         self.radio_description = None
-        hass.data[DATA_ZHA][DATA_ZHA_GATEWAY] = self
         self._log_levels = {
             DEBUG_LEVEL_ORIGINAL: async_capture_log_levels(),
             DEBUG_LEVEL_CURRENT: async_capture_log_levels(),
@@ -122,7 +123,11 @@ class ZHAGateway:
         radio_details = RADIO_TYPES[radio_type]
         radio = radio_details[ZHA_GW_RADIO]()
         self.radio_description = radio_details[ZHA_GW_RADIO_DESCRIPTION]
-        await radio.connect(usb_path, baudrate)
+        try:
+            await radio.connect(usb_path, baudrate)
+        except (SerialException, OSError) as exception:
+            _LOGGER.error("Couldn't open serial port for ZHA: %s", str(exception))
+            raise ConfigEntryNotReady
 
         if CONF_DATABASE in self._config:
             database = self._config[CONF_DATABASE]
@@ -133,7 +138,22 @@ class ZHAGateway:
         apply_application_controller_patch(self)
         self.application_controller.add_listener(self)
         self.application_controller.groups.add_listener(self)
-        await self.application_controller.startup(auto_form=True)
+
+        try:
+            res = await self.application_controller.startup(auto_form=True)
+            if res is False:
+                await self.application_controller.shutdown()
+                raise ConfigEntryNotReady
+        except asyncio.TimeoutError as exception:
+            _LOGGER.error(
+                "Couldn't start %s coordinator",
+                radio_details[ZHA_GW_RADIO_DESCRIPTION],
+                exc_info=exception,
+            )
+            radio.close()
+            raise ConfigEntryNotReady from exception
+
+        self._hass.data[DATA_ZHA][DATA_ZHA_GATEWAY] = self
         self._hass.data[DATA_ZHA][DATA_ZHA_BRIDGE_ID] = str(
             self.application_controller.ieee
         )