Added platform retry support

dcmeglio · Jan 8, 2023 · 6c73744 · 6c73744
1 parent 7ca29ee
commit 6c73744
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 17 deletions.
diff --git a/custom_components/waste_management/manifest.json b/custom_components/waste_management/manifest.json
@@ -1,5 +1,5 @@
 {
-  "version": "0.1.3",
+  "version": "0.1.4",
   "domain": "waste_management",
   "name": "Waste Management",
   "config_flow": true,

diff --git a/custom_components/waste_management/sensor.py b/custom_components/waste_management/sensor.py
@@ -5,6 +5,7 @@
 from .const import CONF_ACCOUNT, CONF_SERVICES
 from homeassistant.components.sensor import SensorEntity
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
+from homeassistant.exceptions import PlatformNotReady
 import homeassistant.helpers.event
 from homeassistant.core import HomeAssistant
 
@@ -17,23 +18,26 @@
 async def async_setup_entry(hass: HomeAssistant, config, add_entities):
     config_data = config.data
     entities = []
-    client = WMClient(config_data[CONF_USERNAME], config_data[CONF_PASSWORD])
-    await client.async_authenticate()
-    await client.async_okta_authorize()
-    wm_services = await client.async_get_services(config_data[CONF_ACCOUNT])
-    for svc_id in config_data[CONF_SERVICES]:
-        name = next(x.name for x in wm_services if x.id == svc_id)
-        entities.append(
-            WasteManagementSensorEntity(
-                hass,
-                name,
-                config_data[CONF_USERNAME],
-                config_data[CONF_PASSWORD],
-                config_data[CONF_ACCOUNT],
-                svc_id,
+    try:
+        client = WMClient(config_data[CONF_USERNAME], config_data[CONF_PASSWORD])
+        await client.async_authenticate()
+        await client.async_okta_authorize()
+        wm_services = await client.async_get_services(config_data[CONF_ACCOUNT])
+        for svc_id in config_data[CONF_SERVICES]:
+            name = next(x.name for x in wm_services if x.id == svc_id)
+            entities.append(
+                WasteManagementSensorEntity(
+                    hass,
+                    name,
+                    config_data[CONF_USERNAME],
+                    config_data[CONF_PASSWORD],
+                    config_data[CONF_ACCOUNT],
+                    svc_id,
+                )
             )
-        )
-    add_entities(entities, True)
+        add_entities(entities, True)
+    except Exception as ex:
+        raise PlatformNotReady("Error setting up sensor platform: %s", ex) from ex
 
 
 class WasteManagementSensorEntity(SensorEntity):