diff --git a/homeassistant/components/vultr/sensor.py b/homeassistant/components/vultr/sensor.py index 01506d4f47e..f4a9055b54c 100644 --- a/homeassistant/components/vultr/sensor.py +++ b/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.""" diff --git a/tests/components/vultr/test_sensor.py b/tests/components/vultr/test_sensor.py index e1dbda1dd04..bacffe8e6af 100644 --- a/tests/components/vultr/test_sensor.py +++ b/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(