core/homeassistant/components/flume/sensor.py

"""Sensor for displaying the number of result from Flume."""

from typing import Any

from pyflume import FlumeAuth, FlumeData
from requests import Session

from homeassistant.components.sensor import (
    SensorDeviceClass,
    SensorEntity,
    SensorEntityDescription,
    SensorStateClass,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import UnitOfVolume
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import StateType

from .const import (
    DEVICE_SCAN_INTERVAL,
    DOMAIN,
    FLUME_AUTH,
    FLUME_DEVICES,
    FLUME_HTTP_SESSION,
    FLUME_TYPE_SENSOR,
    KEY_DEVICE_ID,
    KEY_DEVICE_LOCATION,
    KEY_DEVICE_LOCATION_NAME,
    KEY_DEVICE_LOCATION_TIMEZONE,
    KEY_DEVICE_TYPE,
)
from .coordinator import FlumeDeviceDataUpdateCoordinator
from .entity import FlumeEntity
from .util import get_valid_flume_devices

FLUME_QUERIES_SENSOR: tuple[SensorEntityDescription, ...] = (
    SensorEntityDescription(
        key="current_interval",
        translation_key="current_interval",
        suggested_display_precision=2,
        native_unit_of_measurement=f"{UnitOfVolume.GALLONS}/m",
        state_class=SensorStateClass.MEASUREMENT,
    ),
    SensorEntityDescription(
        key="month_to_date",
        translation_key="month_to_date",
        suggested_display_precision=2,
        native_unit_of_measurement=UnitOfVolume.GALLONS,
        device_class=SensorDeviceClass.WATER,
        state_class=SensorStateClass.TOTAL_INCREASING,
    ),
    SensorEntityDescription(
        key="week_to_date",
        translation_key="week_to_date",
        suggested_display_precision=2,
        native_unit_of_measurement=UnitOfVolume.GALLONS,
        device_class=SensorDeviceClass.WATER,
        state_class=SensorStateClass.TOTAL_INCREASING,
    ),
    SensorEntityDescription(
        key="today",
        translation_key="today",
        suggested_display_precision=2,
        native_unit_of_measurement=UnitOfVolume.GALLONS,
        device_class=SensorDeviceClass.WATER,
        state_class=SensorStateClass.TOTAL_INCREASING,
    ),
    SensorEntityDescription(
        key="last_60_min",
        translation_key="last_60_min",
        suggested_display_precision=2,
        native_unit_of_measurement=f"{UnitOfVolume.GALLONS}/h",
        state_class=SensorStateClass.MEASUREMENT,
    ),
    SensorEntityDescription(
        key="last_24_hrs",
        translation_key="last_24_hrs",
        suggested_display_precision=2,
        native_unit_of_measurement=f"{UnitOfVolume.GALLONS}/d",
        state_class=SensorStateClass.MEASUREMENT,
    ),
    SensorEntityDescription(
        key="last_30_days",
        translation_key="last_30_days",
        suggested_display_precision=2,
        native_unit_of_measurement=f"{UnitOfVolume.GALLONS}/mo",
        state_class=SensorStateClass.MEASUREMENT,
    ),
)


def make_flume_datas(
    http_session: Session, flume_auth: FlumeAuth, flume_devices: list[dict[str, Any]]
) -> dict[str, FlumeData]:
    """Create FlumeData objects for each device."""
    flume_datas: dict[str, FlumeData] = {}
    for device in flume_devices:
        device_id = device[KEY_DEVICE_ID]
        device_timezone = device[KEY_DEVICE_LOCATION][KEY_DEVICE_LOCATION_TIMEZONE]
        flume_data = FlumeData(
            flume_auth,
            device_id,
            device_timezone,
            scan_interval=DEVICE_SCAN_INTERVAL,
            update_on_init=False,
            http_session=http_session,
        )
        flume_datas[device_id] = flume_data
    return flume_datas


async def async_setup_entry(
    hass: HomeAssistant,
    config_entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up the Flume sensor."""

    flume_domain_data = hass.data[DOMAIN][config_entry.entry_id]
    flume_devices = flume_domain_data[FLUME_DEVICES]
    flume_auth: FlumeAuth = flume_domain_data[FLUME_AUTH]
    http_session: Session = flume_domain_data[FLUME_HTTP_SESSION]
    flume_devices = [
        device
        for device in get_valid_flume_devices(flume_devices)
        if device[KEY_DEVICE_TYPE] == FLUME_TYPE_SENSOR
    ]
    flume_entity_list: list[FlumeSensor] = []
    flume_datas = await hass.async_add_executor_job(
        make_flume_datas, http_session, flume_auth, flume_devices
    )

    for device in flume_devices:
        device_id: str = device[KEY_DEVICE_ID]
        device_location_name = device[KEY_DEVICE_LOCATION][KEY_DEVICE_LOCATION_NAME]
        flume_device = flume_datas[device_id]

        coordinator = FlumeDeviceDataUpdateCoordinator(
            hass=hass, flume_device=flume_device
        )

        flume_entity_list.extend(
            [
                FlumeSensor(
                    coordinator=coordinator,
                    description=description,
                    device_id=device_id,
                    location_name=device_location_name,
                )
                for description in FLUME_QUERIES_SENSOR
            ]
        )

    async_add_entities(flume_entity_list)


class FlumeSensor(FlumeEntity[FlumeDeviceDataUpdateCoordinator], SensorEntity):
    """Representation of the Flume sensor."""

    @property
    def native_value(self) -> StateType:
        """Return the state of the sensor."""
        sensor_key = self.entity_description.key
        if sensor_key not in self.coordinator.flume_device.values:
            return None

        return self.coordinator.flume_device.values[sensor_key]