"""Support for Ridwell buttons.""" from typing import Any from aioridwell.errors import RidwellError from aioridwell.model import EventState, RidwellAccount from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from .coordinator import RidwellConfigEntry, RidwellDataUpdateCoordinator from .entity import RidwellEntity SWITCH_DESCRIPTION = SwitchEntityDescription( key="opt_in", translation_key="opt_in", ) async def async_setup_entry( hass: HomeAssistant, entry: RidwellConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up Ridwell sensors based on a config entry.""" coordinator = entry.runtime_data async_add_entities( RidwellSwitch(coordinator, account, SWITCH_DESCRIPTION) for account in coordinator.accounts.values() ) class RidwellSwitch(RidwellEntity, SwitchEntity): """Define a Ridwell switch.""" def __init__( self, coordinator: RidwellDataUpdateCoordinator, account: RidwellAccount, description: SwitchEntityDescription, ) -> None: """Initialize.""" super().__init__(coordinator, account) self._attr_unique_id = f"{account.account_id}_{description.key}" self.entity_description = description @property def is_on(self) -> bool: """Return True if entity is on.""" return self.next_pickup_event.state is EventState.SCHEDULED async def async_turn_off(self, **kwargs: Any) -> None: """Turn the switch off.""" try: await self.next_pickup_event.async_opt_out() except RidwellError as err: raise HomeAssistantError(f"Error while opting out: {err}") from err await self.coordinator.async_request_refresh() async def async_turn_on(self, **kwargs: Any) -> None: """Turn the switch on.""" try: await self.next_pickup_event.async_opt_in() except RidwellError as err: raise HomeAssistantError(f"Error while opting in: {err}") from err await self.coordinator.async_request_refresh()