core/homeassistant/components/lamarzocco/switch.py

"""Switch platform for La Marzocco espresso machines."""

from collections.abc import Callable, Coroutine
from dataclasses import dataclass
from typing import Any

from pylamarzocco.const import BoilerType
from pylamarzocco.exceptions import RequestNotSuccessful
from pylamarzocco.lm_machine import LaMarzoccoMachine
from pylamarzocco.models import LaMarzoccoMachineConfig

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

from .const import DOMAIN
from .coordinator import LaMarzoccoConfigEntry, LaMarzoccoUpdateCoordinator
from .entity import LaMarzoccoBaseEntity, LaMarzoccoEntity, LaMarzoccoEntityDescription


@dataclass(frozen=True, kw_only=True)
class LaMarzoccoSwitchEntityDescription(
    LaMarzoccoEntityDescription,
    SwitchEntityDescription,
):
    """Description of a La Marzocco Switch."""

    control_fn: Callable[[LaMarzoccoMachine, bool], Coroutine[Any, Any, bool]]
    is_on_fn: Callable[[LaMarzoccoMachineConfig], bool]


ENTITIES: tuple[LaMarzoccoSwitchEntityDescription, ...] = (
    LaMarzoccoSwitchEntityDescription(
        key="main",
        translation_key="main",
        name=None,
        control_fn=lambda machine, state: machine.set_power(state),
        is_on_fn=lambda config: config.turned_on,
    ),
    LaMarzoccoSwitchEntityDescription(
        key="steam_boiler_enable",
        translation_key="steam_boiler",
        control_fn=lambda machine, state: machine.set_steam(state),
        is_on_fn=lambda config: config.boilers[BoilerType.STEAM].enabled,
    ),
    LaMarzoccoSwitchEntityDescription(
        key="smart_standby_enabled",
        translation_key="smart_standby_enabled",
        entity_category=EntityCategory.CONFIG,
        control_fn=lambda machine, state: machine.set_smart_standby(
            enabled=state,
            mode=machine.config.smart_standby.mode,
            minutes=machine.config.smart_standby.minutes,
        ),
        is_on_fn=lambda config: config.smart_standby.enabled,
    ),
)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: LaMarzoccoConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up switch entities and services."""

    coordinator = entry.runtime_data

    entities: list[SwitchEntity] = []
    entities.extend(
        LaMarzoccoSwitchEntity(coordinator, description)
        for description in ENTITIES
        if description.supported_fn(coordinator)
    )

    entities.extend(
        LaMarzoccoAutoOnOffSwitchEntity(coordinator, wake_up_sleep_entry_id)
        for wake_up_sleep_entry_id in coordinator.device.config.wake_up_sleep_entries
    )

    async_add_entities(entities)


class LaMarzoccoSwitchEntity(LaMarzoccoEntity, SwitchEntity):
    """Switches representing espresso machine power, prebrew, and auto on/off."""

    entity_description: LaMarzoccoSwitchEntityDescription

    async def async_turn_on(self, **kwargs: Any) -> None:
        """Turn device on."""
        try:
            await self.entity_description.control_fn(self.coordinator.device, True)
        except RequestNotSuccessful as exc:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="switch_on_error",
                translation_placeholders={"key": self.entity_description.key},
            ) from exc
        self.async_write_ha_state()

    async def async_turn_off(self, **kwargs: Any) -> None:
        """Turn device off."""
        try:
            await self.entity_description.control_fn(self.coordinator.device, False)
        except RequestNotSuccessful as exc:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="switch_off_error",
                translation_placeholders={"key": self.entity_description.key},
            ) from exc
        self.async_write_ha_state()

    @property
    def is_on(self) -> bool:
        """Return true if device is on."""
        return self.entity_description.is_on_fn(self.coordinator.device.config)


class LaMarzoccoAutoOnOffSwitchEntity(LaMarzoccoBaseEntity, SwitchEntity):
    """Switch representing espresso machine auto on/off."""

    coordinator: LaMarzoccoUpdateCoordinator
    _attr_translation_key = "auto_on_off"

    def __init__(
        self,
        coordinator: LaMarzoccoUpdateCoordinator,
        identifier: str,
    ) -> None:
        """Initialize the switch."""
        super().__init__(coordinator, f"auto_on_off_{identifier}")
        self._identifier = identifier
        self._attr_translation_placeholders = {"id": identifier}
        self.entity_category = EntityCategory.CONFIG

    async def _async_enable(self, state: bool) -> None:
        """Enable or disable the auto on/off schedule."""
        wake_up_sleep_entry = self.coordinator.device.config.wake_up_sleep_entries[
            self._identifier
        ]
        wake_up_sleep_entry.enabled = state
        try:
            await self.coordinator.device.set_wake_up_sleep(wake_up_sleep_entry)
        except RequestNotSuccessful as exc:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="auto_on_off_error",
                translation_placeholders={"id": self._identifier, "state": str(state)},
            ) from exc
        self.async_write_ha_state()

    async def async_turn_on(self, **kwargs: Any) -> None:
        """Turn switch on."""
        await self._async_enable(True)

    async def async_turn_off(self, **kwargs: Any) -> None:
        """Turn switch off."""
        await self._async_enable(False)

    @property
    def is_on(self) -> bool:
        """Return true if switch is on."""
        return self.coordinator.device.config.wake_up_sleep_entries[
            self._identifier
        ].enabled