core/homeassistant/components/weheat/coordinator.py

"""Define a custom coordinator for the Weheat heatpump integration."""

from datetime import timedelta

from weheat.abstractions.discovery import HeatPumpDiscovery
from weheat.abstractions.heat_pump import HeatPump
from weheat.exceptions import (
    ApiException,
    BadRequestException,
    ForbiddenException,
    NotFoundException,
    ServiceException,
    UnauthorizedException,
)

from homeassistant.const import CONF_ACCESS_TOKEN
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryAuthFailed
from homeassistant.helpers.config_entry_oauth2_flow import OAuth2Session
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed

from .const import API_URL, DOMAIN, LOGGER, UPDATE_INTERVAL

EXCEPTIONS = (
    ServiceException,
    NotFoundException,
    ForbiddenException,
    BadRequestException,
    ApiException,
)


class WeheatDataUpdateCoordinator(DataUpdateCoordinator[HeatPump]):
    """A custom coordinator for the Weheat heatpump integration."""

    def __init__(
        self,
        hass: HomeAssistant,
        session: OAuth2Session,
        heat_pump: HeatPumpDiscovery.HeatPumpInfo,
    ) -> None:
        """Initialize the data coordinator."""
        super().__init__(
            hass,
            logger=LOGGER,
            name=DOMAIN,
            update_interval=timedelta(seconds=UPDATE_INTERVAL),
        )
        self.heat_pump_info = heat_pump
        self._heat_pump_data = HeatPump(API_URL, heat_pump.uuid)

        self.session = session

    @property
    def heatpump_id(self) -> str:
        """Return the heat pump id."""
        return self.heat_pump_info.uuid

    @property
    def readable_name(self) -> str | None:
        """Return the readable name of the heat pump."""
        if self.heat_pump_info.name:
            return self.heat_pump_info.name
        return self.heat_pump_info.model

    @property
    def model(self) -> str:
        """Return the model of the heat pump."""
        return self.heat_pump_info.model

    def fetch_data(self) -> HeatPump:
        """Get the data from the API."""
        try:
            self._heat_pump_data.get_status(self.session.token[CONF_ACCESS_TOKEN])
        except UnauthorizedException as error:
            raise ConfigEntryAuthFailed from error
        except EXCEPTIONS as error:
            raise UpdateFailed(error) from error

        return self._heat_pump_data

    async def _async_update_data(self) -> HeatPump:
        """Fetch data from the API."""
        await self.session.async_ensure_token_valid()

        return await self.hass.async_add_executor_job(self.fetch_data)