From 85d4acb2293e98d47b8615d415c9e4a7f35eef77 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gr=C3=A9goire=20Seux?= Date: Sun, 29 Aug 2021 21:00:35 +0200 Subject: [PATCH] Migrate to SensorEntity base class Goal is to be usable as a source of energy in the energy dashboard. Starting with HA 2021.9 (still in beta), sensors will be have long term stats and be usable in energy dashboard. It also gives a better integration with base classes of home assistant and decrease code in this custom component. Change-Id: Ide1967ac181c0cf7f093cff539aae2b1cdccd9b4 --- custom_components/prometheus_query/sensor.py | 37 ++++++-------------- 1 file changed, 10 insertions(+), 27 deletions(-) diff --git a/custom_components/prometheus_query/sensor.py b/custom_components/prometheus_query/sensor.py index f0a7850..41c883c 100644 --- a/custom_components/prometheus_query/sensor.py +++ b/custom_components/prometheus_query/sensor.py @@ -6,6 +6,7 @@ import homeassistant.helpers.config_validation as cv from datetime import timedelta from homeassistant.helpers.config_validation import PLATFORM_SCHEMA from homeassistant.helpers.entity import Entity +from homeassistant.components.sensor import SensorEntity from homeassistant.util import Throttle from homeassistant.const import ( CONF_NAME, @@ -47,43 +48,24 @@ def setup_platform(hass, config, add_entities, discovery_info=None): 'state_class': str(config.get(CONF_STATE_CLASS)), 'unique_id': str(config.get(CONF_UNIQUE_ID)), } + add_entities([PrometheusQuery(prom_data)], True) -class PrometheusQuery(Entity): - """Representation of a Sensor.""" +class PrometheusQuery(SensorEntity): + """Representation of a Sensor based on Prometheus""" def __init__(self, prom_data): """Initialize the sensor.""" self._url = prom_data["url"] self._query = prom_data["query"] - self._name = prom_data["name"] + self._attr_name = prom_data["name"] self._state = None - self._unit_of_measurement = prom_data["unit"] - self._state_class = prom_data["state_class"] + self._attr_native_unit_of_measurement = prom_data["unit"] + self._attr_state_class = prom_data["state_class"] self._attr_unique_id = f"${prom_data['url']}$${prom_data['query']}" if prom_data["unique_id"] is not None: self._attr_unique_id = prom_data["unique_id"] - @property - def name(self): - """Return the name of the sensor.""" - return self._name - - @property - def state(self): - """Return the state of the sensor.""" - return self._state - - @property - def unit_of_measurement(self): - """Return the unit of measurement.""" - return self._unit_of_measurement - - @property - def state_class(self): - """Return the state_class of the sensor""" - return self._state_class - def update(self): """Fetch new state data for the sensor. This is the only method that should fetch new data for Home Assistant. @@ -92,8 +74,9 @@ class PrometheusQuery(Entity): response = requests.get(self._url, params={'query': self._query}) if (response): results = response.json()['data']['result'] - self._state = results[0]['value'][1] + self._attr_native_value = results[0]['value'][1] else: - self._state = STATE_UNKNOWN + self._attr_native_value = STATE_UNKNOWN + self._attr_state = self._attr_native_value except URLCallError: _LOGGER.error("Error when retrieving update data")