From 677ab5837f02d97896ec35c75193628673b523cd Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Thu, 4 May 2023 19:20:58 +0200
Subject: [PATCH] Add entity descriptions to switch of Rituals Perfume Genie
 (#92531)

---
 .../rituals_perfume_genie/switch.py           | 64 +++++++++++++++----
 1 file changed, 51 insertions(+), 13 deletions(-)

diff --git a/homeassistant/components/rituals_perfume_genie/switch.py b/homeassistant/components/rituals_perfume_genie/switch.py
index 42be84f9278..85e29f7f1ca 100644
--- a/homeassistant/components/rituals_perfume_genie/switch.py
+++ b/homeassistant/components/rituals_perfume_genie/switch.py
@@ -1,9 +1,13 @@
 """Support for Rituals Perfume Genie switches."""
 from __future__ import annotations
 
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
 from typing import Any
 
-from homeassistant.components.switch import SwitchEntity
+from pyrituals import Diffuser
+
+from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -13,6 +17,33 @@ from .coordinator import RitualsDataUpdateCoordinator
 from .entity import DiffuserEntity
 
 
+@dataclass
+class RitualsEntityDescriptionMixin:
+    """Mixin values for Rituals entities."""
+
+    is_on_fn: Callable[[Diffuser], bool]
+    turn_on_fn: Callable[[Diffuser], Awaitable[None]]
+    turn_off_fn: Callable[[Diffuser], Awaitable[None]]
+
+
+@dataclass
+class RitualsSwitchEntityDescription(
+    SwitchEntityDescription, RitualsEntityDescriptionMixin
+):
+    """Class describing Rituals switch entities."""
+
+
+ENTITY_DESCRIPTIONS = (
+    RitualsSwitchEntityDescription(
+        key="is_on",
+        icon="mdi:fan",
+        is_on_fn=lambda diffuser: diffuser.is_on,
+        turn_on_fn=lambda diffuser: diffuser.turn_on(),
+        turn_off_fn=lambda diffuser: diffuser.turn_off(),
+    ),
+)
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
@@ -24,36 +55,43 @@ async def async_setup_entry(
     ]
 
     async_add_entities(
-        DiffuserSwitch(coordinator) for coordinator in coordinators.values()
+        RitualsSwitchEntity(coordinator, description)
+        for coordinator in coordinators.values()
+        for description in ENTITY_DESCRIPTIONS
     )
 
 
-class DiffuserSwitch(DiffuserEntity, SwitchEntity):
+class RitualsSwitchEntity(DiffuserEntity, SwitchEntity):
     """Representation of a diffuser switch."""
 
-    _attr_icon = "mdi:fan"
+    entity_description: RitualsSwitchEntityDescription
 
-    def __init__(self, coordinator: RitualsDataUpdateCoordinator) -> None:
+    def __init__(
+        self,
+        coordinator: RitualsDataUpdateCoordinator,
+        description: RitualsSwitchEntityDescription,
+    ) -> None:
         """Initialize the diffuser switch."""
         super().__init__(coordinator)
-        self._attr_is_on = self.coordinator.diffuser.is_on
-        self._attr_unique_id = f"{coordinator.diffuser.hublot}-is_on"
+        self.entity_description = description
+        self._attr_unique_id = f"{coordinator.diffuser.hublot}-{description.key}"
         self._attr_name = coordinator.diffuser.name
+        self._attr_is_on = description.is_on_fn(coordinator.diffuser)
 
     async def async_turn_on(self, **kwargs: Any) -> None:
-        """Turn the device on."""
-        await self.coordinator.diffuser.turn_on()
+        """Turn the switch on."""
+        await self.entity_description.turn_on_fn(self.coordinator.diffuser)
         self._attr_is_on = True
         self.async_write_ha_state()
 
     async def async_turn_off(self, **kwargs: Any) -> None:
-        """Turn the device off."""
-        await self.coordinator.diffuser.turn_off()
+        """Turn the switch off."""
+        await self.entity_description.turn_off_fn(self.coordinator.diffuser)
         self._attr_is_on = False
         self.async_write_ha_state()
 
     @callback
     def _handle_coordinator_update(self) -> None:
         """Handle updated data from the coordinator."""
-        self._attr_is_on = self.coordinator.diffuser.is_on
-        self.async_write_ha_state()
+        self._attr_is_on = self.entity_description.is_on_fn(self.coordinator.diffuser)
+        super()._handle_coordinator_update()