Fix forward_entry_setup error, fix asyncio error (probably)

custom_components/eia_hourly_demand/__init__.py

@@ -1,4 +1,5 @@
 """The EIA Energy Data component."""
+
 from .const import DOMAIN
 
 CONF_API_KEY = "api_key"
@@ -13,16 +14,20 @@ def setup(hass, config):
 async def async_setup_entry(hass, entry):
     """Set up EIA Energy Data from a config entry."""
     hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = entry.data[CONF_API_KEY]
+    hass.data[DOMAIN][entry.entry_id] = entry.data
 
-    hass.async_create_task(
-        hass.config_entries.async_forward_entry_setup(entry, "sensor")
-    )
+    # Forward the setup to the sensor platform
+    await hass.config_entries.async_forward_entry_setup(entry, "sensor")
 
     return True
 
 
 async def async_unload_entry(hass, entry):
     """Unload the EIA Energy sensor platform."""
     await hass.config_entries.async_forward_entry_unload(entry, "sensor")
+
+    # Clean up the entry from hass.data
+    if entry.entry_id in hass.data[DOMAIN]:
+        hass.data[DOMAIN].pop(entry.entry_id)
+
     return True

custom_components/eia_hourly_demand/config_flow.py

@@ -1,4 +1,5 @@
-""" Config Flow for EIA Integration """
+"""Config Flow for EIA Integration"""
+
 from datetime import timedelta, date
 import aiohttp
 import voluptuous as vol
@@ -17,13 +18,19 @@
 
 
 class EIAConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
+    """Handle a config flow for EIA integration."""
+
+    VERSION = 1
+
     async def async_step_user(self, user_input=None):
+        """Handle the initial step."""
         errors = {}
 
         if user_input is not None:
             # Validate the user input
-            if not await self._validate_input(user_input):
-                errors["base"] = "invalid API Key or Balancing Authority"
+            valid = await self._validate_input(user_input)
+            if not valid:
+                errors["base"] = "invalid_api_key_or_ba_id"
             else:
                 # If validation is successful, create and return the config entry
                 return self.async_create_entry(title=user_input[BA_ID], data=user_input)
@@ -44,21 +51,28 @@ async def _validate_input(self, user_input):
         api_key = user_input[API_KEY]
         ba_id = user_input[BA_ID]
 
-        # Check if the API key is valid by making a test API call
+        # Check if the API key and BA ID are valid by making a test API call
         start_date = (date.today() - timedelta(days=7)).strftime("%Y-%m-%d")
         url = EIA_URL.format(api_key=api_key, ba_id=ba_id, start_date=start_date)
 
         async with aiohttp.ClientSession() as session:
             try:
-                timeout = aiohttp.ClientTimeout(total=5)
+                timeout = aiohttp.ClientTimeout(
+                    total=10
+                )  # Adjusted timeout for reliability
                 async with session.get(url, timeout=timeout) as response:
                     if response.status != 200:
                         return False
                     data = await response.json()
-# Not needed                   
-#                    if data["response"]["total"] == 0:
-#                        return False
+                    if not data["response"]["data"]:
+                        return False
             except aiohttp.ClientConnectorError:
                 return False
+            except aiohttp.ClientResponseError as e:
+                return False
+            except aiohttp.TimeoutError:
+                return False
+            except Exception as e:
+                return False
 
         return True

custom_components/eia_hourly_demand/manifest.json

@@ -8,5 +8,5 @@
     "iot_class":  "cloud_polling",
     "issue_tracker": "https://github.com/hokiebrian/eia_hourly_demand/issues",
     "requirements": [],
-    "version": "1.0.13"
+    "version": "1.0.14"
   }

custom_components/eia_hourly_demand/sensor.py

@@ -1,10 +1,14 @@
-""" Setup EIA Sensor """
+"""
+Setup EIA Sensor
+"""
+
 import logging
 from datetime import timedelta, date
-import json
 import aiohttp
 from homeassistant.core import HomeAssistant
 from homeassistant.components.sensor import SensorEntity, SensorStateClass
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
@@ -22,7 +26,12 @@
 )
 
 
-async def async_setup_entry(hass: HomeAssistant, config_entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+):
+    """Set up the EIA sensor entry."""
     api_key = config_entry.data[API_KEY]
     ba_id = config_entry.data[BA_ID]
     eia_data = hass.data[DOMAIN][config_entry.entry_id]
@@ -31,54 +40,62 @@ async def async_setup_entry(hass: HomeAssistant, config_entry, async_add_entitie
 
 
 class EIASensor(SensorEntity):
+    """Representation of an EIA Sensor."""
+
     _attr_icon = "mdi:factory"
     _attr_native_unit_of_measurement = "MWh"
     _attr_state_class: SensorStateClass = SensorStateClass.MEASUREMENT
 
-    def __init__(self, api_key, ba_id, eia_data):
+    def __init__(self, api_key: str, ba_id: str, eia_data: dict):
+        """Initialize the sensor."""
         self._api_key = api_key
         self._ba_id = ba_id
         self._eia_data = eia_data
         self._state = None
 
     @property
-    def name(self):
+    def name(self) -> str:
+        """Return the name of the sensor."""
         return f"Hourly Demand {self._ba_id}"
 
     @property
-    def state(self):
+    def state(self) -> float:
+        """Return the state of the sensor."""
         return self._state
 
     @property
-    def unique_id(self):
+    def unique_id(self) -> str:
+        """Return a unique ID for the sensor."""
         return f"HourlyMWh{self._ba_id}"
 
-    async def async_update(self):
+    async def async_update(self) -> None:
+        """Fetch new state data for the sensor."""
         start_date = (date.today() - timedelta(days=7)).strftime("%Y-%m-%d")
         url = EIA_URL.format(
             api_key=self._api_key, ba_id=self._ba_id, start_date=start_date
         )
-        _LOGGER.debug(f"Data {url}")
+        _LOGGER.debug(f"Fetching data from URL: {url}")
+
         async with aiohttp.ClientSession() as session:
             try:
-                timeout = aiohttp.ClientTimeout(total=5)
+                timeout = aiohttp.ClientTimeout(total=10)
                 async with session.get(url, timeout=timeout) as response:
+                    response.raise_for_status()  # Raise an error for bad HTTP status codes
                     data = await response.json()
-                    if data["response"]["data"][0]["value"] is None:
-                        value_as_float = 0
-                    else:
-                        value_as_float = float(data["response"]["data"][0]["value"])
-                    self._state = value_as_float
+                    self._state = float(data["response"]["data"][0]["value"])
             except aiohttp.ClientConnectorError as e:
-                _LOGGER.debug(f"Connection Error: {e}")
+                _LOGGER.error(f"Connection Error: {e}")
                 self._state = None
             except (IndexError, KeyError) as e:
-                _LOGGER.error("Data Error, no data returned")
-                _LOGGER.debug(f"Data Error: {e}")
+                _LOGGER.error("Data Error: Invalid or no data returned")
+                _LOGGER.debug(f"Error details: {e}")
+                self._state = None
+            except aiohttp.TimeoutError as e:
+                _LOGGER.error("Timeout Error: Request timed out")
+                _LOGGER.debug(f"Error details: {e}")
                 self._state = None
-            except asyncio.TimeoutError as e:
-                _LOGGER.error("Timeout Error, asyncio")
-                _LOGGER.debug(f"Data Error: {e}")
+            except aiohttp.ClientResponseError as e:
+                _LOGGER.error(f"HTTP Error: {e.status} - {e.message}")
                 self._state = None
             except Exception as e:
                 _LOGGER.error(f"An unexpected error occurred: {e}")