diff --git a/custom_components/idrac_power/__init__.py b/custom_components/idrac_power/__init__.py index e1f1e53..a63b508 100644 --- a/custom_components/idrac_power/__init__.py +++ b/custom_components/idrac_power/__init__.py @@ -3,32 +3,36 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant +from homeassistant.exceptions import ConfigEntryNotReady from .const import DOMAIN, DATA_IDRAC_REST_CLIENT, HOST, USERNAME, PASSWORD, CONF_INTERVAL, CONF_INTERVAL_DEFAULT -from .idrac_rest import IdracMock, IdracRest +from .idrac_rest import IdracMock, IdracRest, CannotConnect async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up the iDrac connection from a config entry.""" - if entry.data[HOST] == 'MOCK': - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = { - DATA_IDRAC_REST_CLIENT: IdracMock( - entry.data[HOST], - entry.data[USERNAME], - entry.data[PASSWORD], - entry.data.get(CONF_INTERVAL, CONF_INTERVAL_DEFAULT) - ) - } - else: - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = { - DATA_IDRAC_REST_CLIENT: IdracRest( - entry.data[HOST], - entry.data[USERNAME], - entry.data[PASSWORD], - entry.data.get(CONF_INTERVAL, CONF_INTERVAL_DEFAULT) - ) - } + try: + if entry.data[HOST] == 'MOCK': + hass.data.setdefault(DOMAIN, {})[entry.entry_id] = { + DATA_IDRAC_REST_CLIENT: IdracMock( + entry.data[HOST], + entry.data[USERNAME], + entry.data[PASSWORD], + entry.data.get(CONF_INTERVAL, CONF_INTERVAL_DEFAULT) + ) + } + else: + hass.data.setdefault(DOMAIN, {})[entry.entry_id] = { + DATA_IDRAC_REST_CLIENT: IdracRest( + entry.data[HOST], + entry.data[USERNAME], + entry.data[PASSWORD], + entry.data.get(CONF_INTERVAL, CONF_INTERVAL_DEFAULT) + ) + } + except CannotConnect: + raise ConfigEntryNotReady hass.async_create_task( hass.config_entries.async_forward_entry_setup(