Add a switch to opt-in to/opt-out of the next Ridwell pickup (#62293)

* Add buttons to opt into/out of the next Ridwell pickup

* Buttons finished

* Coverage

* better name

* Move to switch

* Clean up

* Coverage

* Use correct exception

.coveragerc

@@ -887,6 +887,7 @@ omit =
     homeassistant/components/rest/switch.py
     homeassistant/components/ridwell/__init__.py
     homeassistant/components/ridwell/sensor.py
+    homeassistant/components/ridwell/switch.py
     homeassistant/components/ring/camera.py
     homeassistant/components/ripple/sensor.py
     homeassistant/components/rocketchat/notify.py

homeassistant/components/ridwell/__init__.py

@@ -12,14 +12,25 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
-from homeassistant.helpers import aiohttp_client
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+from homeassistant.helpers import aiohttp_client, entity_registry as er
+from homeassistant.helpers.entity import EntityDescription
+from homeassistant.helpers.update_coordinator import (
+    CoordinatorEntity,
+    DataUpdateCoordinator,
+    UpdateFailed,
+)
 
-from .const import DATA_ACCOUNT, DATA_COORDINATOR, DOMAIN, LOGGER
+from .const import (
+    DATA_ACCOUNT,
+    DATA_COORDINATOR,
+    DOMAIN,
+    LOGGER,
+    SENSOR_TYPE_NEXT_PICKUP,
+)
 
 DEFAULT_UPDATE_INTERVAL = timedelta(hours=1)
 
-PLATFORMS: list[Platform] = [Platform.SENSOR]
+PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.SWITCH]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
@@ -82,3 +93,43 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
+
+
+async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Migrate an old config entry."""
+    version = entry.version
+
+    LOGGER.debug("Migrating from version %s", version)
+
+    # 1 -> 2: Update unique ID of existing, single sensor entity to be consistent with
+    # common format for platforms going forward:
+    if version == 1:
+        version = entry.version = 2
+
+        ent_reg = er.async_get(hass)
+        [entity_entry] = [
+            e for e in ent_reg.entities.values() if e.config_entry_id == entry.entry_id
+        ]
+        new_unique_id = f"{entity_entry.unique_id}_{SENSOR_TYPE_NEXT_PICKUP}"
+        ent_reg.async_update_entity(entity_entry.entity_id, new_unique_id=new_unique_id)
+
+    LOGGER.info("Migration to version %s successful", version)
+
+    return True
+
+
+class RidwellEntity(CoordinatorEntity):
+    """Define a base Ridwell entity."""
+
+    def __init__(
+        self,
+        coordinator: DataUpdateCoordinator,
+        account: RidwellAccount,
+        description: EntityDescription,
+    ) -> None:
+        """Initialize the sensor."""
+        super().__init__(coordinator)
+
+        self._account = account
+        self._attr_unique_id = f"{account.account_id}_{description.key}"
+        self.entity_description = description

homeassistant/components/ridwell/config_flow.py

@@ -32,7 +32,7 @@ STEP_USER_DATA_SCHEMA = vol.Schema(
 class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle a config flow for WattTime."""
 
-    VERSION = 1
+    VERSION = 2
 
     def __init__(self) -> None:
         """Initialize."""

homeassistant/components/ridwell/const.py

@@ -7,3 +7,5 @@ LOGGER = logging.getLogger(__package__)
 
 DATA_ACCOUNT = "account"
 DATA_COORDINATOR = "coordinator"
+
+SENSOR_TYPE_NEXT_PICKUP = "next_pickup"

homeassistant/components/ridwell/sensor.py

@@ -7,49 +7,60 @@ from typing import Any
 
 from aioridwell.model import RidwellAccount, RidwellPickupEvent
 
-from homeassistant.components.sensor import SensorDeviceClass, SensorEntity
+from homeassistant.components.sensor import (
+    SensorDeviceClass,
+    SensorEntity,
+    SensorEntityDescription,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .const import DATA_ACCOUNT, DATA_COORDINATOR, DOMAIN
+from . import RidwellEntity
+from .const import DATA_ACCOUNT, DATA_COORDINATOR, DOMAIN, SENSOR_TYPE_NEXT_PICKUP
 
 ATTR_CATEGORY = "category"
 ATTR_PICKUP_STATE = "pickup_state"
 ATTR_PICKUP_TYPES = "pickup_types"
 ATTR_QUANTITY = "quantity"
 
+SENSOR_DESCRIPTION = SensorEntityDescription(
+    key=SENSOR_TYPE_NEXT_PICKUP,
+    name="Ridwell Pickup",
+    device_class=SensorDeviceClass.DATE,
+)
+
 
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
-    """Set up WattTime sensors based on a config entry."""
+    """Set up Ridwell sensors based on a config entry."""
     accounts = hass.data[DOMAIN][entry.entry_id][DATA_ACCOUNT]
     coordinator = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR]
+
     async_add_entities(
-        [RidwellSensor(coordinator, account) for account in accounts.values()]
+        [
+            RidwellSensor(coordinator, account, SENSOR_DESCRIPTION)
+            for account in accounts.values()
+        ]
     )
 
 
-class RidwellSensor(CoordinatorEntity, SensorEntity):
+class RidwellSensor(RidwellEntity, SensorEntity):
     """Define a Ridwell pickup sensor."""
 
-    _attr_device_class = SensorDeviceClass.DATE
-
     def __init__(
-        self, coordinator: DataUpdateCoordinator, account: RidwellAccount
+        self,
+        coordinator: DataUpdateCoordinator,
+        account: RidwellAccount,
+        description: SensorEntityDescription,
     ) -> None:
-        """Initialize the sensor."""
-        super().__init__(coordinator)
+        """Initialize."""
+        super().__init__(coordinator, account, description)
 
-        self._account = account
-        self._attr_name = f"Ridwell Pickup ({account.address['street1']})"
-        self._attr_unique_id = account.account_id
+        self._attr_name = f"{description.name} ({account.address['street1']})"
 
     @property
     def extra_state_attributes(self) -> Mapping[str, Any]:

homeassistant/components/ridwell/switch.py

@@ -0,0 +1,71 @@
+"""Support for Ridwell buttons."""
+from __future__ import annotations
+
+from typing import Any
+
+from aioridwell.errors import RidwellError
+from aioridwell.model import EventState, RidwellPickupEvent
+
+from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from . import RidwellEntity
+from .const import DATA_ACCOUNT, DATA_COORDINATOR, DOMAIN
+
+SWITCH_TYPE_OPT_IN = "opt_in"
+
+SWITCH_DESCRIPTION = SwitchEntityDescription(
+    key=SWITCH_TYPE_OPT_IN,
+    name="Opt-In to Next Pickup",
+    icon="mdi:calendar-check",
+)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+) -> None:
+    """Set up Ridwell sensors based on a config entry."""
+    accounts = hass.data[DOMAIN][entry.entry_id][DATA_ACCOUNT]
+    coordinator = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR]
+
+    async_add_entities(
+        [
+            RidwellSwitch(coordinator, account, SWITCH_DESCRIPTION)
+            for account in accounts.values()
+        ]
+    )
+
+
+class RidwellSwitch(RidwellEntity, SwitchEntity):
+    """Define a Ridwell button."""
+
+    @property
+    def is_on(self) -> bool:
+        """Return True if entity is on."""
+        event: RidwellPickupEvent = self.coordinator.data[self._account.account_id]
+        return event.state == EventState.SCHEDULED
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn the switch off."""
+        event: RidwellPickupEvent = self.coordinator.data[self._account.account_id]
+
+        try:
+            await 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."""
+        event: RidwellPickupEvent = self.coordinator.data[self._account.account_id]
+
+        try:
+            await event.async_opt_in()
+        except RidwellError as err:
+            raise HomeAssistantError(f"Error while opting in: {err}") from err
+
+        await self.coordinator.async_request_refresh()