Add extra sensors by adding support to the extended ESIOS API with API Token

homeassistant/components/pvpc_hourly_pricing/__init__.py

@@ -1,84 +1,61 @@
 """The pvpc_hourly_pricing integration to collect Spain official electric prices."""
-from collections.abc import Mapping
-from datetime import datetime, timedelta
+from __future__ import annotations
+
+from datetime import timedelta
 import logging
 
-from aiopvpc import DEFAULT_POWER_KW, TARIFFS, PVPCData
-import voluptuous as vol
+from aiopvpc import (
+    DEFAULT_POWER_KW,
+    BadApiTokenAuthError,
+    EsiosApiData,
+    PVPCData,
+    get_enabled_sensor_keys,
+)
 
-from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
-from homeassistant.const import CONF_NAME
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_API_TOKEN
 from homeassistant.core import HomeAssistant, callback
+from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
-import homeassistant.helpers.config_validation as cv
-from homeassistant.helpers.entity_registry import (
-    EntityRegistry,
-    async_get,
-    async_migrate_entries,
-)
-from homeassistant.helpers.typing import ConfigType
+import homeassistant.helpers.entity_registry as er
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 from homeassistant.util import dt as dt_util
 
 from .const import (
     ATTR_POWER,
     ATTR_POWER_P3,
     ATTR_TARIFF,
-    DEFAULT_NAME,
+    CONF_USE_API_TOKEN,
+    DEFAULT_TARIFF,
     DOMAIN,
     PLATFORMS,
 )
 
 _LOGGER = logging.getLogger(__name__)
-_DEFAULT_TARIFF = TARIFFS[0]
-VALID_POWER = vol.All(vol.Coerce(float), vol.Range(min=1.0, max=15.0))
-VALID_TARIFF = vol.In(TARIFFS)
-UI_CONFIG_SCHEMA = vol.Schema(
-    {
-        vol.Required(CONF_NAME, default=DEFAULT_NAME): str,
-        vol.Required(ATTR_TARIFF, default=_DEFAULT_TARIFF): VALID_TARIFF,
-        vol.Required(ATTR_POWER, default=DEFAULT_POWER_KW): VALID_POWER,
-        vol.Required(ATTR_POWER_P3, default=DEFAULT_POWER_KW): VALID_POWER,
-    }
-)
-CONFIG_SCHEMA = vol.Schema(
-    vol.All(cv.deprecated(DOMAIN), {DOMAIN: cv.ensure_list(UI_CONFIG_SCHEMA)}),
-    extra=vol.ALLOW_EXTRA,
-)
-
-
-async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
-    """Set up the electricity price sensor from configuration.yaml."""
-    for conf in config.get(DOMAIN, []):
-        hass.async_create_task(
-            hass.config_entries.flow.async_init(
-                DOMAIN, data=conf, context={"source": SOURCE_IMPORT}
-            )
-        )
-
-    return True
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up pvpc hourly pricing from a config entry."""
+    entity_registry = er.async_get(hass)
     if len(entry.data) == 2:
         defaults = {
-            ATTR_TARIFF: _DEFAULT_TARIFF,
+            ATTR_TARIFF: DEFAULT_TARIFF,
             ATTR_POWER: DEFAULT_POWER_KW,
             ATTR_POWER_P3: DEFAULT_POWER_KW,
+            CONF_API_TOKEN: None,
         }
         data = {**entry.data, **defaults}
         hass.config_entries.async_update_entry(
-            entry, unique_id=_DEFAULT_TARIFF, data=data, options=defaults
+            entry, unique_id=DEFAULT_TARIFF, data=data, options=defaults
         )
 
         @callback
         def update_unique_id(reg_entry):
             """Change unique id for sensor entity, pointing to new tariff."""
-            return {"new_unique_id": _DEFAULT_TARIFF}
+            return {"new_unique_id": DEFAULT_TARIFF}
 
         try:
-            await async_migrate_entries(hass, entry.entry_id, update_unique_id)
+            await er.async_migrate_entries(hass, entry.entry_id, update_unique_id)
             _LOGGER.warning(
                 (
                     "Migrating PVPC sensor from old tariff '%s' to new '%s'. "
@@ -87,15 +64,14 @@ def update_unique_id(reg_entry):
                     "if that is your case"
                 ),
                 entry.data[ATTR_TARIFF],
-                _DEFAULT_TARIFF,
+                DEFAULT_TARIFF,
             )
         except ValueError:
             # there were multiple sensors (with different old tariffs, up to 3),
             # so we leave just one and remove the others
-            ent_reg: EntityRegistry = async_get(hass)
-            for entity_id, reg_entry in ent_reg.entities.items():
+            for entity_id, reg_entry in entity_registry.entities.items():
                 if reg_entry.config_entry_id == entry.entry_id:
-                    ent_reg.async_remove(entity_id)
+                    entity_registry.async_remove(entity_id)
                     _LOGGER.warning(
                         (
                             "Old PVPC Sensor %s is removed "
@@ -108,7 +84,12 @@ def update_unique_id(reg_entry):
             await hass.config_entries.async_remove(entry.entry_id)
             return False
 
-    coordinator = ElecPricesDataUpdateCoordinator(hass, entry)
+    entries = er.async_entries_for_config_entry(entity_registry, entry.entry_id)
+    sensor_keys = get_enabled_sensor_keys(
+        using_private_api=entry.data.get(CONF_USE_API_TOKEN, False),
+        disabled_sensor_ids=[sensor.unique_id for sensor in entries if sensor.disabled],
+    )
+    coordinator = ElecPricesDataUpdateCoordinator(hass, entry, sensor_keys)
     await coordinator.async_config_entry_first_refresh()
 
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
@@ -121,7 +102,7 @@ async def async_update_options(hass: HomeAssistant, entry: ConfigEntry) -> None:
     """Handle options update."""
     if any(
         entry.data.get(attrib) != entry.options.get(attrib)
-        for attrib in (ATTR_TARIFF, ATTR_POWER, ATTR_POWER_P3)
+        for attrib in (ATTR_POWER, ATTR_POWER_P3, CONF_USE_API_TOKEN, CONF_API_TOKEN)
     ):
         # update entry replacing data with new options
         hass.config_entries.async_update_entry(
@@ -138,17 +119,24 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return unload_ok
 
 
-class ElecPricesDataUpdateCoordinator(DataUpdateCoordinator[Mapping[datetime, float]]):
+class ElecPricesDataUpdateCoordinator(DataUpdateCoordinator[EsiosApiData]):
     """Class to manage fetching Electricity prices data from API."""
 
-    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
+    def __init__(
+        self, hass: HomeAssistant, entry: ConfigEntry, sensor_keys: set[str]
+    ) -> None:
         """Initialize."""
+        api_token = None
+        if entry.data.get(CONF_USE_API_TOKEN, False):
+            api_token = entry.data.get(CONF_API_TOKEN)
         self.api = PVPCData(
             session=async_get_clientsession(hass),
             tariff=entry.data[ATTR_TARIFF],
             local_timezone=hass.config.time_zone,
             power=entry.data[ATTR_POWER],
             power_valley=entry.data[ATTR_POWER_P3],
+            api_token=api_token,
+            sensor_keys=tuple(sensor_keys),
         )
         super().__init__(
             hass, _LOGGER, name=DOMAIN, update_interval=timedelta(minutes=30)
@@ -160,11 +148,16 @@ def entry_id(self) -> str:
         """Return entry ID."""
         return self._entry.entry_id
 
-    async def _async_update_data(self) -> Mapping[datetime, float]:
+    async def _async_update_data(self) -> EsiosApiData:
         """Update electricity prices from the ESIOS API."""
-        prices = await self.api.async_update_prices(dt_util.utcnow())
-        self.api.process_state_and_attributes(dt_util.utcnow())
-        if not prices:
+        try:
+            api_data = await self.api.async_update_all(self.data, dt_util.utcnow())
+        except BadApiTokenAuthError as exc:
+            raise ConfigEntryAuthFailed from exc
+        if (
+            not api_data
+            or not api_data.sensors
+            or not all(api_data.availability.values())
+        ):
             raise UpdateFailed
-
-        return prices
+        return api_data