core/homeassistant/components/guardian/switch.py

"""Switches for the Elexa Guardian integration."""
from __future__ import annotations

from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
from typing import Any

from aioguardian import Client
from aioguardian.errors import GuardianError

from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity_platform import AddEntitiesCallback

from . import GuardianData, ValveControllerEntity, ValveControllerEntityDescription
from .const import API_VALVE_STATUS, API_WIFI_STATUS, DOMAIN

ATTR_AVG_CURRENT = "average_current"
ATTR_CONNECTED_CLIENTS = "connected_clients"
ATTR_INST_CURRENT = "instantaneous_current"
ATTR_INST_CURRENT_DDT = "instantaneous_current_ddt"
ATTR_STATION_CONNECTED = "station_connected"
ATTR_TRAVEL_COUNT = "travel_count"

SWITCH_KIND_ONBOARD_AP = "onboard_ap"
SWITCH_KIND_VALVE = "valve"

ON_STATES = {
    "start_opening",
    "opening",
    "finish_opening",
    "opened",
}


@dataclass(frozen=True, kw_only=True)
class ValveControllerSwitchDescription(
    SwitchEntityDescription, ValveControllerEntityDescription
):
    """Describe a Guardian valve controller switch."""

    extra_state_attributes_fn: Callable[[dict[str, Any]], Mapping[str, Any]]
    is_on_fn: Callable[[dict[str, Any]], bool]
    off_fn: Callable[[Client], Awaitable]
    on_fn: Callable[[Client], Awaitable]


async def _async_disable_ap(client: Client) -> None:
    """Disable the onboard AP."""
    await client.wifi.disable_ap()


async def _async_enable_ap(client: Client) -> None:
    """Enable the onboard AP."""
    await client.wifi.enable_ap()


async def _async_close_valve(client: Client) -> None:
    """Close the valve."""
    await client.valve.close()


async def _async_open_valve(client: Client) -> None:
    """Open the valve."""
    await client.valve.open()


VALVE_CONTROLLER_DESCRIPTIONS = (
    ValveControllerSwitchDescription(
        key=SWITCH_KIND_ONBOARD_AP,
        translation_key="onboard_access_point",
        icon="mdi:wifi",
        entity_category=EntityCategory.CONFIG,
        extra_state_attributes_fn=lambda data: {
            ATTR_CONNECTED_CLIENTS: data.get("ap_clients"),
            ATTR_STATION_CONNECTED: data["station_connected"],
        },
        api_category=API_WIFI_STATUS,
        is_on_fn=lambda data: data["ap_enabled"],
        off_fn=_async_disable_ap,
        on_fn=_async_enable_ap,
    ),
    ValveControllerSwitchDescription(
        key=SWITCH_KIND_VALVE,
        translation_key="valve_controller",
        icon="mdi:water",
        api_category=API_VALVE_STATUS,
        extra_state_attributes_fn=lambda data: {
            ATTR_AVG_CURRENT: data["average_current"],
            ATTR_INST_CURRENT: data["instantaneous_current"],
            ATTR_INST_CURRENT_DDT: data["instantaneous_current_ddt"],
            ATTR_TRAVEL_COUNT: data["travel_count"],
        },
        is_on_fn=lambda data: data["state"] in ON_STATES,
        off_fn=_async_close_valve,
        on_fn=_async_open_valve,
    ),
)


async def async_setup_entry(
    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
) -> None:
    """Set up Guardian switches based on a config entry."""
    data: GuardianData = hass.data[DOMAIN][entry.entry_id]

    async_add_entities(
        ValveControllerSwitch(entry, data, description)
        for description in VALVE_CONTROLLER_DESCRIPTIONS
    )


class ValveControllerSwitch(ValveControllerEntity, SwitchEntity):
    """Define a switch related to a Guardian valve controller."""

    entity_description: ValveControllerSwitchDescription

    def __init__(
        self,
        entry: ConfigEntry,
        data: GuardianData,
        description: ValveControllerSwitchDescription,
    ) -> None:
        """Initialize."""
        super().__init__(entry, data.valve_controller_coordinators, description)

        self._client = data.client

    @property
    def extra_state_attributes(self) -> Mapping[str, Any]:
        """Return entity specific state attributes."""
        return self.entity_description.extra_state_attributes_fn(self.coordinator.data)

    @property
    def is_on(self) -> bool:
        """Return True if entity is on."""
        return self.entity_description.is_on_fn(self.coordinator.data)

    async def async_turn_off(self, **kwargs: Any) -> None:
        """Turn the switch off."""
        if not self._attr_is_on:
            return

        try:
            async with self._client:
                await self.entity_description.off_fn(self._client)
        except GuardianError as err:
            raise HomeAssistantError(
                f'Error while turning "{self.entity_id}" off: {err}'
            ) from err

        self._attr_is_on = False
        self.async_write_ha_state()

    async def async_turn_on(self, **kwargs: Any) -> None:
        """Turn the switch on."""
        if self._attr_is_on:
            return

        try:
            async with self._client:
                await self.entity_description.on_fn(self._client)
        except GuardianError as err:
            raise HomeAssistantError(
                f'Error while turning "{self.entity_id}" on: {err}'
            ) from err

        self._attr_is_on = True
        self.async_write_ha_state()