"""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()