Add ConfigEntryNotReady exception

cbrherms · May 6, 2024 · d65bc3b · d65bc3b
1 parent e1931b9
commit d65bc3b
Showing 1 changed file with 23 additions and 19 deletions.
diff --git 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(