Use EntityDescription - ombi (home-assistant#55086)

jborup · Aug 23, 2021 · 354dbe9 · 354dbe9
1 parent dae4053
commit 354dbe9
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 49 deletions.
diff --git a/homeassistant/components/ombi/const.py b/homeassistant/components/ombi/const.py
@@ -1,4 +1,8 @@
 """Support for Ombi."""
+from __future__ import annotations
+
+from homeassistant.components.sensor import SensorEntityDescription
+
 ATTR_SEASON = "season"
 
 CONF_URLBASE = "urlbase"
@@ -13,11 +17,35 @@
 SERVICE_MUSIC_REQUEST = "submit_music_request"
 SERVICE_TV_REQUEST = "submit_tv_request"
 
-SENSOR_TYPES = {
-    "movies": {"type": "Movie requests", "icon": "mdi:movie"},
-    "tv": {"type": "TV show requests", "icon": "mdi:television-classic"},
-    "music": {"type": "Music album requests", "icon": "mdi:album"},
-    "pending": {"type": "Pending requests", "icon": "mdi:clock-alert-outline"},
-    "approved": {"type": "Approved requests", "icon": "mdi:check"},
-    "available": {"type": "Available requests", "icon": "mdi:download"},
-}
+SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
+    SensorEntityDescription(
+        key="movies",
+        name="Movie requests",
+        icon="mdi:movie",
+    ),
+    SensorEntityDescription(
+        key="tv",
+        name="TV show requests",
+        icon="mdi:television-classic",
+    ),
+    SensorEntityDescription(
+        key="music",
+        name="Music album requests",
+        icon="mdi:album",
+    ),
+    SensorEntityDescription(
+        key="pending",
+        name="Pending requests",
+        icon="mdi:clock-alert-outline",
+    ),
+    SensorEntityDescription(
+        key="approved",
+        name="Approved requests",
+        icon="mdi:check",
+    ),
+    SensorEntityDescription(
+        key="available",
+        name="Available requests",
+        icon="mdi:download",
+    ),
+)
diff --git a/homeassistant/components/ombi/sensor.py b/homeassistant/components/ombi/sensor.py
@@ -4,7 +4,7 @@
 
 from pyombi import OmbiError
 
-from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
 
 from .const import DOMAIN, SENSOR_TYPES
 
@@ -18,60 +18,39 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
     if discovery_info is None:
         return
 
-    sensors = []
-
     ombi = hass.data[DOMAIN]["instance"]
 
-    for sensor, sensor_val in SENSOR_TYPES.items():
-        sensor_label = sensor
-        sensor_type = sensor_val["type"]
-        sensor_icon = sensor_val["icon"]
-        sensors.append(OmbiSensor(sensor_label, sensor_type, ombi, sensor_icon))
+    entities = [OmbiSensor(ombi, description) for description in SENSOR_TYPES]
 
-    add_entities(sensors, True)
+    add_entities(entities, True)
 
 
 class OmbiSensor(SensorEntity):
     """Representation of an Ombi sensor."""
 
-    def __init__(self, label, sensor_type, ombi, icon):
+    def __init__(self, ombi, description: SensorEntityDescription):
         """Initialize the sensor."""
-        self._state = None
-        self._label = label
-        self._type = sensor_type
+        self.entity_description = description
         self._ombi = ombi
-        self._icon = icon
-
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return f"Ombi {self._type}"
-
-    @property
-    def icon(self):
-        """Return the icon to use in the frontend."""
-        return self._icon
 
-    @property
-    def native_value(self):
-        """Return the state of the sensor."""
-        return self._state
+        self._attr_name = f"Ombi {description.name}"
 
     def update(self):
         """Update the sensor."""
         try:
-            if self._label == "movies":
-                self._state = self._ombi.movie_requests
-            elif self._label == "tv":
-                self._state = self._ombi.tv_requests
-            elif self._label == "music":
-                self._state = self._ombi.music_requests
-            elif self._label == "pending":
-                self._state = self._ombi.total_requests["pending"]
-            elif self._label == "approved":
-                self._state = self._ombi.total_requests["approved"]
-            elif self._label == "available":
-                self._state = self._ombi.total_requests["available"]
+            sensor_type = self.entity_description.key
+            if sensor_type == "movies":
+                self._attr_native_value = self._ombi.movie_requests
+            elif sensor_type == "tv":
+                self._attr_native_value = self._ombi.tv_requests
+            elif sensor_type == "music":
+                self._attr_native_value = self._ombi.music_requests
+            elif sensor_type == "pending":
+                self._attr_native_value = self._ombi.total_requests["pending"]
+            elif sensor_type == "approved":
+                self._attr_native_value = self._ombi.total_requests["approved"]
+            elif sensor_type == "available":
+                self._attr_native_value = self._ombi.total_requests["available"]
         except OmbiError as err:
             _LOGGER.warning("Unable to update Ombi sensor: %s", err)
-            self._state = None
+            self._attr_native_value = None