fix: use prometheus metrics in hass integration

.vscode/launch.json

@@ -4,6 +4,14 @@
     // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
     "version": "0.2.0",
     "configurations": [
+        {
+            "name": "Python: Current File",
+            "type": "python",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal",
+            "justMyCode": true
+        },
         {
             "name": "Python: Module",
             "type": "python",

custom_components/estimenergy/__init__.py

@@ -2,21 +2,24 @@
 
 import logging
 import asyncio
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.config_entries import ConfigEntry, ConfigType
 from homeassistant.core import HomeAssistant
-from estimenergy.const import SENSOR_TYPES
 
 from .const import PLATFORM, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
 
-async def async_setup(hass, config) -> bool:
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+    """Set up the EstimEnergy component."""
+
     hass.data.setdefault(DOMAIN, {})
     return True
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Set up EstimEnergy from a config entry."""
+
     hass.data.setdefault(DOMAIN, {})
     hass_data = dict(entry.data)
     hass.data[DOMAIN][entry.entry_id] = hass_data
@@ -41,7 +44,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         )
     )
 
-    # Remove config entry from domain.
     if unload_ok:
         hass.data[DOMAIN].pop(entry.entry_id)
 

custom_components/estimenergy/const.py

@@ -4,6 +4,5 @@
 PLATFORM = Platform.SENSOR
 DOMAIN = "estimenergy"
 
-CONF_NAME = "collector_name"
 CONF_HOST = "host"
 CONF_PORT = "port"

custom_components/estimenergy/coordinator.py

@@ -2,61 +2,56 @@
 
 import logging
 from datetime import timedelta
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_HOST, CONF_PORT
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+from prometheus_client.parser import text_string_to_metric_families
 
-from estimenergy_client import EstimEnergyClient
-
+from estimenergy.client import EstimEnergyClient
+from estimenergy.const import METRICS, Metric
 
 _LOGGER = logging.getLogger(__name__)
 
 
 class EstimEnergyCoordinator(DataUpdateCoordinator):
     """Data coordinator for EstimEnergy."""
 
-    def __init__(self, hass: HomeAssistant, host: str, port: int) -> None:
+    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
         super().__init__(
             hass,
             _LOGGER,
             update_interval=timedelta(seconds=5),
         )
 
         self.hass = hass
-        self.host = host
-        self.port = port
-        self.client = EstimEnergyClient(host, port)
-        self.collector_id = None
-        self._data = None
-        self.collector = None
-
-    async def initialize(self) -> None:
-        """Initialize the EstimEnergy API connection."""
-
-        collectors = await self.hass.async_add_executor_job(self.client.get_collectors)
+        self.host = entry.data[CONF_HOST]
+        self.port = entry.data[CONF_PORT]
+        self.client = EstimEnergyClient(self.host, self.port)
 
-        for collector in collectors:
-            if collector["name"] == self.name:
-                self.collector_id = collector["id"]
-                break
+    def get_value_for_metric(self, metric: Metric, samples: list):
+        """Get the value for a metric."""
 
-        if not self.collector_id:
-            raise CollectorNotFoundError(self.name)
+        for sample in samples:
+            if sample.name == metric.json_key:
+                return sample.value
+        return None
 
     async def _async_update_data(self):
         """Refresh data from API."""
 
-        if self.collector_id is None:
-            return None
-
-        self.collector = await self.hass.async_add_executor_job(
-            self.client.get_collector, self.collector_id
-        )
-
-        return self.collector
-
-
-class CollectorNotFoundError(Exception):
-    """Custom Exception for a collector not being found."""
-
-    def __init__(self, collector_name: str) -> None:
-        super().__init__(f"Collector with name {collector_name} not found!")
+        metrics_text = await self.hass.async_add_executor_job(self.client.get_metrics)
+        metrics = text_string_to_metric_families(metrics_text)
+        samples = [
+            sample
+            for family in text_string_to_metric_families(metrics)
+            for sample in family.samples
+            if family.name in [metric.json_key for metric in METRICS]
+        ]
+
+        return {
+            name: {
+                metric: self.get_value_for_metric(metric, samples) for metric in METRICS
+            }
+            for name in set([sample.labels["name"] for sample in samples])
+        }

custom_components/estimenergy/sensor.py

@@ -11,25 +11,11 @@
 from homeassistant.components.sensor import SensorEntity
 from homeassistant.const import CURRENCY_EURO, PERCENTAGE
 from homeassistant.const import UnitOfEnergy
-from estimenergy.const import (
-    JSON_BILLING_MONTH,
-    JSON_DATA,
-    JSON_DAY_KWH,
-    JSON_DAY_COST,
-    JSON_DAY_COST_DIFFERENCE,
-    JSON_MONTH_KWH_RAW,
-    JSON_YEAR_KWH_RAW,
-    JSON_MONTH_KWH,
-    JSON_MONTH_COST,
-    JSON_MONTH_COST_DIFFERENCE,
-    JSON_YEAR_KWH,
-    SENSOR_TYPE_FRIENDLY_NAME,
-    SENSOR_TYPE_JSON,
-    SENSOR_TYPE_UNIQUE_ID,
-    CONF_NAME,
-    CONF_HOST,
-    CONF_PORT,
-)
+from estimenergy.client import EstimEnergyClient
+from estimenergy.const import METRICS, Metric, MetricType, MetricPeriod
+
+from .const import CONF_HOST, CONF_PORT
+from .coordinator import EstimEnergyCoordinator
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -78,13 +64,7 @@ def __init__(
     @property
     def device_class(self) -> SensorDeviceClass | None:
         """Return the class of this entity."""
-        if self.json_key in [
-            JSON_DAY_KWH,
-            JSON_MONTH_KWH_RAW,
-            JSON_YEAR_KWH_RAW,
-            JSON_MONTH_KWH,
-            JSON_YEAR_KWH,
-        ]:
+        if self.metric.metric_type == MetricType.ENERGY:
             return SensorDeviceClass.ENERGY
 
         if self.metric.metric_type in [MetricType.COST, MetricType.COST_DIFFERENCE]:
@@ -116,7 +96,7 @@ def last_reset(self) -> datetime | None:
         if self.metric.metric_period == MetricPeriod.TOTAL:
             return None
 
-        billing_month = self.coordinator.data[JSON_BILLING_MONTH]
+        billing_month = self.collector["billing_month"]
         now = datetime.now()
 
         return now.replace(
@@ -134,12 +114,12 @@ def native_value(self) -> int | None:
         """Return the state of the sensor."""
         if (
             self.coordinator.data is None
-            or JSON_METRICS not in self.coordinator.data
-            or self.metric.json_key not in self.coordinator.data[JSON_METRICS]
+            or "metrics" not in self.coordinator.data
+            or self.metric.json_key not in self.coordinator.data["metrics"]
         ):
             return None
 
-        return self.coordinator.data[JSON_METRICS][self.metric.json_key]
+        return self.coordinator.data[self.collector["name"]][self.metric.json_key]
 
     @property
     def suggested_display_precision(self) -> int | None:
@@ -149,13 +129,7 @@ def suggested_display_precision(self) -> int | None:
     @property
     def native_unit_of_measurement(self) -> str | None:
         """Return the unit of measurement of the sensor."""
-        if self.json_key in [
-            JSON_DAY_KWH,
-            JSON_MONTH_KWH_RAW,
-            JSON_YEAR_KWH_RAW,
-            JSON_MONTH_KWH,
-            JSON_YEAR_KWH,
-        ]:
+        if self.device_class == SensorDeviceClass.ENERGY:
             return UnitOfEnergy.KILO_WATT_HOUR
 
         if self.device_class == SensorDeviceClass.MONETARY:

estimenergy/client/client.py

@@ -26,3 +26,8 @@ async def async_get_collectors(self):
         url = f"http://{self.host}:{self.port}/collector"
         response = await requests.get(url)
         return response.json()
+
+    def get_metrics(self):
+        url = f"http://{self.host}:{self.port}/metrics"
+        response = requests.get(url)
+        return response.text

estimenergy/const.py

@@ -30,15 +30,15 @@ def __init__(self, metric_type: MetricType, metric_period: MetricPeriod, is_pred
 
     @property
     def json_key(self) -> str:
-        return f"{self.metric_period.value[0]}_{self.metric_type.value[0]}{'_predicted' if self.is_predicted else ''}{'_raw' if self.is_raw else ''}"
+        return f"estimenergy_{self.metric_period.value[0]}_{self.metric_type.value[0]}{'_predicted' if self.is_predicted else ''}{'_raw' if self.is_raw else ''}"
 
     @property
     def friendly_name(self) -> str:
         return f"{self.metric_period.value[1]} {self.metric_type.value[1]} {'(Predicted)' if self.is_predicted else ''} {'(Raw)' if self.is_raw else ''}"
 
     def create_gauge(self) -> Gauge:
         return Gauge(
-            f"estimenergy_{self.json_key}",
+            f"{self.json_key}",
             f"EstimEnergy {self.friendly_name}",
             ["name", "id"],
         )