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

@@ -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):
+class PrometheusQuery(SensorEntity):
-    """Representation of a Sensor."""
+    """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._attr_native_unit_of_measurement = prom_data["unit"]
-        self._state_class = prom_data["state_class"]
+        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")