core/homeassistant/components/osoenergy/__init__.py

"""Support for the OSO Energy devices and services."""

from typing import Any, Generic, TypeVar

from aiohttp.web_exceptions import HTTPException
from apyosoenergyapi import OSOEnergy
from apyosoenergyapi.helper.const import (
    OSOEnergyBinarySensorData,
    OSOEnergySensorData,
    OSOEnergyWaterHeaterData,
)
from apyosoenergyapi.helper.osoenergy_exceptions import OSOEnergyReauthRequired

from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_API_KEY, Platform
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
from homeassistant.helpers import aiohttp_client
from homeassistant.helpers.device_registry import DeviceInfo
from homeassistant.helpers.entity import Entity

from .const import DOMAIN

_OSOEnergyT = TypeVar(
    "_OSOEnergyT",
    OSOEnergyBinarySensorData,
    OSOEnergySensorData,
    OSOEnergyWaterHeaterData,
)

MANUFACTURER = "OSO Energy"
PLATFORMS = [
    Platform.SENSOR,
    Platform.WATER_HEATER,
]
PLATFORM_LOOKUP = {
    Platform.SENSOR: "sensor",
    Platform.WATER_HEATER: "water_heater",
}


async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Set up OSO Energy from a config entry."""
    subscription_key = entry.data[CONF_API_KEY]
    websession = aiohttp_client.async_get_clientsession(hass)
    osoenergy = OSOEnergy(subscription_key, websession)

    osoenergy_config = dict(entry.data)

    hass.data.setdefault(DOMAIN, {})

    try:
        devices: Any = await osoenergy.session.start_session(osoenergy_config)
    except HTTPException as error:
        raise ConfigEntryNotReady from error
    except OSOEnergyReauthRequired as err:
        raise ConfigEntryAuthFailed from err

    hass.data[DOMAIN][entry.entry_id] = osoenergy

    platforms = set()
    for ha_type, oso_type in PLATFORM_LOOKUP.items():
        device_list = devices.get(oso_type, [])
        if device_list:
            platforms.add(ha_type)
    if platforms:
        await hass.config_entries.async_forward_entry_setups(entry, platforms)
    return True


async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
    """Unload a config entry."""
    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
    if unload_ok:
        hass.data[DOMAIN].pop(entry.entry_id)

    return unload_ok


class OSOEnergyEntity(Entity, Generic[_OSOEnergyT]):
    """Initiate OSO Energy Base Class."""

    _attr_has_entity_name = True

    def __init__(self, osoenergy: OSOEnergy, entity_data: _OSOEnergyT) -> None:
        """Initialize the instance."""
        self.osoenergy = osoenergy
        self.entity_data = entity_data
        self._attr_device_info = DeviceInfo(
            identifiers={(DOMAIN, entity_data.device_id)},
            manufacturer=MANUFACTURER,
            model=entity_data.device_type,
            name=entity_data.device_name,
        )