Refactor ZHA setup (#32959)

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

* Cleanup.
This commit is contained in:
Alexei Chetroi 2020-03-19 12:37:47 -04:00 committed by GitHub
parent d33a3ca90f
commit c8c81d493d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 39 additions and 19 deletions

View File

@ -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, {})
hass.data[DATA_ZHA][DATA_ZHA_DISPATCHERS] = []
hass.data[DATA_ZHA][DATA_ZHA_PLATFORM_LOADED] = asyncio.Event()
zha_data = hass.data.setdefault(DATA_ZHA, {})
config = zha_data.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()
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 hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].async_update_device_storage()
await zha_data[DATA_ZHA_GATEWAY].shutdown()
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())

View File

@ -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
)