diff --git a/homeassistant/components/ovo_energy/__init__.py b/homeassistant/components/ovo_energy/__init__.py index 436180407f4..ec5d1c7cafa 100644 --- a/homeassistant/components/ovo_energy/__init__.py +++ b/homeassistant/components/ovo_energy/__init__.py @@ -2,23 +2,19 @@ from __future__ import annotations -import asyncio -from datetime import timedelta import logging import aiohttp from ovoenergy import OVOEnergy -from ovoenergy.models import OVODailyUsage from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.helpers.aiohttp_client import async_get_clientsession -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed -from homeassistant.util import dt as dt_util from .const import CONF_ACCOUNT, DATA_CLIENT, DATA_COORDINATOR, DOMAIN +from .coordinator import OVOEnergyDataUpdateCoordinator _LOGGER = logging.getLogger(__name__) @@ -47,33 +43,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: _LOGGER.warning(exception) raise ConfigEntryNotReady from exception - async def async_update_data() -> OVODailyUsage: - """Fetch data from OVO Energy.""" - if (custom_account := entry.data.get(CONF_ACCOUNT)) is not None: - client.custom_account_id = custom_account - - async with asyncio.timeout(10): - try: - authenticated = await client.authenticate( - entry.data[CONF_USERNAME], - entry.data[CONF_PASSWORD], - ) - except aiohttp.ClientError as exception: - raise UpdateFailed(exception) from exception - if not authenticated: - raise ConfigEntryAuthFailed("Not authenticated with OVO Energy") - return await client.get_daily_usage(dt_util.utcnow().strftime("%Y-%m")) - - coordinator = DataUpdateCoordinator[OVODailyUsage]( - hass, - _LOGGER, - config_entry=entry, - # Name of the data. For logging purposes. - name="sensor", - update_method=async_update_data, - # Polling interval. Will only be polled if there are subscribers. - update_interval=timedelta(seconds=3600), - ) + coordinator = OVOEnergyDataUpdateCoordinator(hass, entry, client) hass.data.setdefault(DOMAIN, {}) hass.data[DOMAIN][entry.entry_id] = { diff --git a/homeassistant/components/ovo_energy/coordinator.py b/homeassistant/components/ovo_energy/coordinator.py new file mode 100644 index 00000000000..6d06fd56092 --- /dev/null +++ b/homeassistant/components/ovo_energy/coordinator.py @@ -0,0 +1,61 @@ +"""Coordinator for the OVO Energy integration.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta +import logging + +import aiohttp +from ovoenergy import OVOEnergy +from ovoenergy.models import OVODailyUsage + +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import CONF_PASSWORD, CONF_USERNAME +from homeassistant.core import HomeAssistant +from homeassistant.exceptions import ConfigEntryAuthFailed +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed +from homeassistant.util import dt as dt_util + +from .const import CONF_ACCOUNT + +_LOGGER = logging.getLogger(__name__) + + +class OVOEnergyDataUpdateCoordinator(DataUpdateCoordinator[OVODailyUsage]): + """Class to manage fetching OVO Energy data.""" + + config_entry: ConfigEntry + + def __init__( + self, + hass: HomeAssistant, + config_entry: ConfigEntry, + client: OVOEnergy, + ) -> None: + """Initialize.""" + super().__init__( + hass, + _LOGGER, + config_entry=config_entry, + name="sensor", + update_interval=timedelta(seconds=3600), + ) + self.client = client + + async def _async_update_data(self) -> OVODailyUsage: + """Fetch data from OVO Energy.""" + if (custom_account := self.config_entry.data.get(CONF_ACCOUNT)) is not None: + self.client.custom_account_id = custom_account + + async with asyncio.timeout(10): + try: + authenticated = await self.client.authenticate( + self.config_entry.data[CONF_USERNAME], + self.config_entry.data[CONF_PASSWORD], + ) + except aiohttp.ClientError as exception: + raise UpdateFailed(exception) from exception + if not authenticated: + raise ConfigEntryAuthFailed("Not authenticated with OVO Energy") + return await self.client.get_daily_usage(dt_util.utcnow().strftime("%Y-%m")) diff --git a/homeassistant/components/ovo_energy/entity.py b/homeassistant/components/ovo_energy/entity.py index ed8a24b0542..1839f0bae4c 100644 --- a/homeassistant/components/ovo_energy/entity.py +++ b/homeassistant/components/ovo_energy/entity.py @@ -3,25 +3,22 @@ from __future__ import annotations from ovoenergy import OVOEnergy -from ovoenergy.models import OVODailyUsage from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo -from homeassistant.helpers.update_coordinator import ( - CoordinatorEntity, - DataUpdateCoordinator, -) +from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import DOMAIN +from .coordinator import OVOEnergyDataUpdateCoordinator -class OVOEnergyEntity(CoordinatorEntity[DataUpdateCoordinator[OVODailyUsage]]): +class OVOEnergyEntity(CoordinatorEntity[OVOEnergyDataUpdateCoordinator]): """Defines a base OVO Energy entity.""" _attr_has_entity_name = True def __init__( self, - coordinator: DataUpdateCoordinator[OVODailyUsage], + coordinator: OVOEnergyDataUpdateCoordinator, client: OVOEnergy, ) -> None: """Initialize the OVO Energy entity.""" diff --git a/homeassistant/components/ovo_energy/sensor.py b/homeassistant/components/ovo_energy/sensor.py index e2ac9410cbc..a42e193e4b5 100644 --- a/homeassistant/components/ovo_energy/sensor.py +++ b/homeassistant/components/ovo_energy/sensor.py @@ -21,10 +21,10 @@ from homeassistant.const import UnitOfEnergy from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.typing import StateType -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from homeassistant.util import dt as dt_util from .const import DATA_CLIENT, DATA_COORDINATOR, DOMAIN +from .coordinator import OVOEnergyDataUpdateCoordinator from .entity import OVOEnergyDeviceEntity SCAN_INTERVAL = timedelta(seconds=300) @@ -118,9 +118,9 @@ async def async_setup_entry( async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up OVO Energy sensor based on a config entry.""" - coordinator: DataUpdateCoordinator[OVODailyUsage] = hass.data[DOMAIN][ - entry.entry_id - ][DATA_COORDINATOR] + coordinator: OVOEnergyDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][ + DATA_COORDINATOR + ] client: OVOEnergy = hass.data[DOMAIN][entry.entry_id][DATA_CLIENT] entities = [] @@ -161,12 +161,11 @@ async def async_setup_entry( class OVOEnergySensor(OVOEnergyDeviceEntity, SensorEntity): """Define a OVO Energy sensor.""" - coordinator: DataUpdateCoordinator[DataUpdateCoordinator[OVODailyUsage]] entity_description: OVOEnergySensorEntityDescription def __init__( self, - coordinator: DataUpdateCoordinator[OVODailyUsage], + coordinator: OVOEnergyDataUpdateCoordinator, description: OVOEnergySensorEntityDescription, client: OVOEnergy, ) -> None: