Use EntityDescription - vultr (#55789)

homeassistant/components/vultr/sensor.py

@@ -1,9 +1,15 @@
 """Support for monitoring the state of Vultr Subscriptions."""
+from __future__ import annotations
+
 import logging
 
 import voluptuous as vol
 
-from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
+from homeassistant.components.sensor import (
+    PLATFORM_SCHEMA,
+    SensorEntity,
+    SensorEntityDescription,
+)
 from homeassistant.const import CONF_MONITORED_CONDITIONS, CONF_NAME, DATA_GIGABYTES
 import homeassistant.helpers.config_validation as cv
 
@@ -17,22 +23,29 @@ from . import (
 _LOGGER = logging.getLogger(__name__)
 
 DEFAULT_NAME = "Vultr {} {}"
-MONITORED_CONDITIONS = {
-    ATTR_CURRENT_BANDWIDTH_USED: [
-        "Current Bandwidth Used",
-        DATA_GIGABYTES,
-        "mdi:chart-histogram",
-    ],
-    ATTR_PENDING_CHARGES: ["Pending Charges", "US$", "mdi:currency-usd"],
-}
+SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
+    SensorEntityDescription(
+        key=ATTR_CURRENT_BANDWIDTH_USED,
+        name="Current Bandwidth Used",
+        native_unit_of_measurement=DATA_GIGABYTES,
+        icon="mdi:chart-histogram",
+    ),
+    SensorEntityDescription(
+        key=ATTR_PENDING_CHARGES,
+        name="Pending Charges",
+        native_unit_of_measurement="US$",
+        icon="mdi:currency-usd",
+    ),
+)
+SENSOR_KEYS: list[str] = [desc.key for desc in SENSOR_TYPES]
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
     {
         vol.Required(CONF_SUBSCRIPTION): cv.string,
         vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
-        vol.Optional(
-            CONF_MONITORED_CONDITIONS, default=list(MONITORED_CONDITIONS)
-        ): vol.All(cv.ensure_list, [vol.In(MONITORED_CONDITIONS)]),
+        vol.Optional(CONF_MONITORED_CONDITIONS, default=SENSOR_KEYS): vol.All(
+            cv.ensure_list, [vol.In(SENSOR_KEYS)]
+        ),
     }
 )
 
@@ -41,68 +54,55 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
     """Set up the Vultr subscription (server) sensor."""
     vultr = hass.data[DATA_VULTR]
 
-    subscription = config.get(CONF_SUBSCRIPTION)
-    name = config.get(CONF_NAME)
-    monitored_conditions = config.get(CONF_MONITORED_CONDITIONS)
+    subscription = config[CONF_SUBSCRIPTION]
+    name = config[CONF_NAME]
+    monitored_conditions = config[CONF_MONITORED_CONDITIONS]
 
     if subscription not in vultr.data:
         _LOGGER.error("Subscription %s not found", subscription)
         return
 
-    sensors = []
+    entities = [
+        VultrSensor(vultr, subscription, name, description)
+        for description in SENSOR_TYPES
+        if description.key in monitored_conditions
+    ]
 
-    for condition in monitored_conditions:
-        sensors.append(VultrSensor(vultr, subscription, condition, name))
-
-    add_entities(sensors, True)
+    add_entities(entities, True)
 
 
 class VultrSensor(SensorEntity):
     """Representation of a Vultr subscription sensor."""
 
-    def __init__(self, vultr, subscription, condition, name):
+    def __init__(self, vultr, subscription, name, description: SensorEntityDescription):
         """Initialize a new Vultr sensor."""
+        self.entity_description = description
         self._vultr = vultr
-        self._condition = condition
         self._name = name
 
         self.subscription = subscription
         self.data = None
 
-        condition_info = MONITORED_CONDITIONS[condition]
-
-        self._condition_name = condition_info[0]
-        self._units = condition_info[1]
-        self._icon = condition_info[2]
-
     @property
     def name(self):
         """Return the name of the sensor."""
         try:
-            return self._name.format(self._condition_name)
+            return self._name.format(self.entity_description.name)
         except IndexError:
             try:
-                return self._name.format(self.data["label"], self._condition_name)
+                return self._name.format(
+                    self.data["label"], self.entity_description.name
+                )
             except (KeyError, TypeError):
                 return self._name
 
-    @property
-    def icon(self):
-        """Return the icon used in the frontend if any."""
-        return self._icon
-
-    @property
-    def native_unit_of_measurement(self):
-        """Return the unit of measurement to present the value in."""
-        return self._units
-
     @property
     def native_value(self):
         """Return the value of this given sensor type."""
         try:
-            return round(float(self.data.get(self._condition)), 2)
+            return round(float(self.data.get(self.entity_description.key)), 2)
         except (TypeError, ValueError):
-            return self.data.get(self._condition)
+            return self.data.get(self.entity_description.key)
 
     def update(self):
         """Update state of sensor."""

tests/components/vultr/test_sensor.py

@@ -39,12 +39,12 @@ class TestVultrSensorSetup(unittest.TestCase):
             {
                 CONF_NAME: vultr.DEFAULT_NAME,
                 CONF_SUBSCRIPTION: "576965",
-                CONF_MONITORED_CONDITIONS: vultr.MONITORED_CONDITIONS,
+                CONF_MONITORED_CONDITIONS: vultr.SENSOR_KEYS,
             },
             {
                 CONF_NAME: "Server {}",
                 CONF_SUBSCRIPTION: "123456",
-                CONF_MONITORED_CONDITIONS: vultr.MONITORED_CONDITIONS,
+                CONF_MONITORED_CONDITIONS: vultr.SENSOR_KEYS,
             },
             {
                 CONF_NAME: "VPS Charges",
@@ -126,7 +126,7 @@ class TestVultrSensorSetup(unittest.TestCase):
             vultr.PLATFORM_SCHEMA(
                 {
                     CONF_PLATFORM: base_vultr.DOMAIN,
-                    CONF_MONITORED_CONDITIONS: vultr.MONITORED_CONDITIONS,
+                    CONF_MONITORED_CONDITIONS: vultr.SENSOR_KEYS,
                 }
             )
         with pytest.raises(vol.Invalid):  # Bad monitored_conditions
@@ -154,7 +154,9 @@ class TestVultrSensorSetup(unittest.TestCase):
             base_vultr.setup(self.hass, VALID_CONFIG)
 
         bad_conf = {
-            CONF_MONITORED_CONDITIONS: vultr.MONITORED_CONDITIONS
+            CONF_NAME: "Vultr {} {}",
+            CONF_SUBSCRIPTION: "",
+            CONF_MONITORED_CONDITIONS: vultr.SENSOR_KEYS,
         }  # No subs at all
 
         no_sub_setup = vultr.setup_platform(