Use EntityDescription - vultr (#55789)

This commit is contained in:
Marc Mueller 2021-09-06 09:44:33 +02:00 committed by GitHub
parent cc6a0d2f8d
commit 96db04213b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 48 additions and 46 deletions

View File

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

View File

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