From 328ba6f762059cd8cab229e89ebf0176e6a45606 Mon Sep 17 00:00:00 2001
From: joostlek <joostlek@outlook.com>
Date: Thu, 24 Aug 2023 12:47:14 +0200
Subject: [PATCH] Use entity descriptions in Switcher

---
 .../components/switcher_kis/sensor.py         | 75 +++++++------------
 1 file changed, 29 insertions(+), 46 deletions(-)

diff --git a/homeassistant/components/switcher_kis/sensor.py b/homeassistant/components/switcher_kis/sensor.py
index 2c74f14cb5cb41..0b6263a6b2ef58 100644
--- a/homeassistant/components/switcher_kis/sensor.py
+++ b/homeassistant/components/switcher_kis/sensor.py
@@ -1,19 +1,19 @@
 """Switcher integration Sensor platform."""
 from __future__ import annotations
 
-from dataclasses import dataclass
-
 from aioswitcher.device import DeviceCategory
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
     SensorEntity,
+    SensorEntityDescription,
     SensorStateClass,
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfElectricCurrent, UnitOfPower
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
@@ -22,48 +22,36 @@
 from . import SwitcherDataUpdateCoordinator
 from .const import SIGNAL_DEVICE_ADD
 
-
-@dataclass
-class AttributeDescription:
-    """Class to describe a sensor."""
-
-    name: str
-    icon: str | None = None
-    unit: str | None = None
-    device_class: SensorDeviceClass | None = None
-    state_class: SensorStateClass | None = None
-    default_enabled: bool = True
-
-
-POWER_SENSORS = {
-    "power_consumption": AttributeDescription(
+POWER_SENSORS: list[SensorEntityDescription] = [
+    SensorEntityDescription(
+        key="power_consumption",
         name="Power Consumption",
-        unit=UnitOfPower.WATT,
+        native_unit_of_measurement=UnitOfPower.WATT,
         device_class=SensorDeviceClass.POWER,
         state_class=SensorStateClass.MEASUREMENT,
     ),
-    "electric_current": AttributeDescription(
+    SensorEntityDescription(
+        key="electric_current",
         name="Electric Current",
-        unit=UnitOfElectricCurrent.AMPERE,
+        native_unit_of_measurement=UnitOfElectricCurrent.AMPERE,
         device_class=SensorDeviceClass.CURRENT,
         state_class=SensorStateClass.MEASUREMENT,
     ),
-}
-
-TIME_SENSORS = {
-    "remaining_time": AttributeDescription(
-        name="Remaining Time",
-        icon="mdi:av-timer",
+]
+TIME_SENSORS: list[SensorEntityDescription] = [
+    SensorEntityDescription(
+        key="remaining_time", name="Remaining Time", icon="mdi:av-timer"
     ),
-    "auto_off_set": AttributeDescription(
+    SensorEntityDescription(
+        key="auto_off_set",
         name="Auto Shutdown",
         icon="mdi:progress-clock",
-        default_enabled=False,
+        entity_registry_enabled_default=False,
     ),
-}
+]
 
 POWER_PLUG_SENSORS = POWER_SENSORS
-WATER_HEATER_SENSORS = {**POWER_SENSORS, **TIME_SENSORS}
+WATER_HEATER_SENSORS = [*POWER_SENSORS, *TIME_SENSORS]
 
 
 async def async_setup_entry(
@@ -78,13 +66,13 @@ def async_add_sensors(coordinator: SwitcherDataUpdateCoordinator) -> None:
         """Add sensors from Switcher device."""
         if coordinator.data.device_type.category == DeviceCategory.POWER_PLUG:
             async_add_entities(
-                SwitcherSensorEntity(coordinator, attribute, info)
-                for attribute, info in POWER_PLUG_SENSORS.items()
+                SwitcherSensorEntity(coordinator, description)
+                for description in POWER_PLUG_SENSORS
             )
         elif coordinator.data.device_type.category == DeviceCategory.WATER_HEATER:
             async_add_entities(
-                SwitcherSensorEntity(coordinator, attribute, info)
-                for attribute, info in WATER_HEATER_SENSORS.items()
+                SwitcherSensorEntity(coordinator, description)
+                for description in WATER_HEATER_SENSORS
             )
 
     config_entry.async_on_unload(
@@ -100,28 +88,23 @@ class SwitcherSensorEntity(
     def __init__(
         self,
         coordinator: SwitcherDataUpdateCoordinator,
-        attribute: str,
-        description: AttributeDescription,
+        description: SensorEntityDescription,
     ) -> None:
         """Initialize the entity."""
         super().__init__(coordinator)
-        self.attribute = attribute
+        self.entity_description = description
 
         # Entity class attributes
         self._attr_name = f"{coordinator.name} {description.name}"
-        self._attr_icon = description.icon
-        self._attr_native_unit_of_measurement = description.unit
-        self._attr_device_class = description.device_class
-        self._attr_entity_registry_enabled_default = description.default_enabled
 
         self._attr_unique_id = (
-            f"{coordinator.device_id}-{coordinator.mac_address}-{attribute}"
+            f"{coordinator.device_id}-{coordinator.mac_address}-{description.key}"
+        )
+        self._attr_device_info = DeviceInfo(
+            connections={(dr.CONNECTION_NETWORK_MAC, coordinator.mac_address)}
         )
-        self._attr_device_info = {
-            "connections": {(dr.CONNECTION_NETWORK_MAC, coordinator.mac_address)}
-        }
 
     @property
     def native_value(self) -> StateType:
         """Return value of sensor."""
-        return getattr(self.coordinator.data, self.attribute)  # type: ignore[no-any-return]
+        return getattr(self.coordinator.data, self.entity_description.key)  # type: ignore[no-any-return]