From 4fbe9c5660290af70f4b2b52caa9211e8d1f0a5c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gr=C3=A9goire=20Seux?= <kamaradclimber@gmail.com>
Date: Thu, 26 Aug 2021 23:04:05 +0200
Subject: [PATCH] Support state class propoerty for sensor

Fix lfasci/homeassistant-prometheus-query#5

Goal is to make a first pass to allow to be exposed in the energy
dashboard.
This requires usage of HA version 2021.9 (which is beta at time of
writing) if we want to allow total_increase state_class
---
 README.md                                    |  6 ++++++
 custom_components/prometheus_query/sensor.py | 13 +++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/README.md b/README.md
index c49cebd..4412a85 100644
--- a/README.md
+++ b/README.md
@@ -13,6 +13,7 @@ sensor:
     prometheus_url: http://localhost:9090
     prometheus_query: temperature{location="Pisa",province="PI",region="Tuscany"}
     unit_of_measurement: "°C"
+    state_class: total_increasing
 ```
 
 ### Configuration Variables
@@ -32,5 +33,10 @@ sensor:
 -  unit_of_measurement
 
   (string)(Optional) Defines the unit of measurement of the sensor, if any.
+  
+- state_class
+
+  (string)(Optional) Defines the type of sensor. `measurement` for metrics that are gauges,
+                     `total_increasing` for metrics that are counters.
 
 It's a custom component so it must be downloaded under /custom_components folder.
diff --git a/custom_components/prometheus_query/sensor.py b/custom_components/prometheus_query/sensor.py
index 98f4318..70aa2a7 100644
--- a/custom_components/prometheus_query/sensor.py
+++ b/custom_components/prometheus_query/sensor.py
@@ -12,11 +12,16 @@ from homeassistant.const import (
     CONF_UNIT_OF_MEASUREMENT,
     STATE_UNKNOWN,
 )
+from homeassistant.components.sensor import (
+    DEVICE_CLASSES_SCHEMA,
+    STATE_CLASSES_SCHEMA,
+)
 
 from prometheus_client import Summary
 
 CONF_PROMETHEUS_URL = 'prometheus_url'
 CONF_PROMETHEUS_QUERY = 'prometheus_query'
+CONF_STATE_CLASS = 'state_class'
 SCAN_INTERVAL = timedelta(seconds=600)
 
 _LOGGER = logging.getLogger(__name__)
@@ -27,6 +32,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Required(CONF_PROMETHEUS_QUERY): cv.string,
     vol.Required(CONF_NAME): cv.string,
     vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
+    vol.Optional(CONF_STATE_CLASS): STATE_CLASSES_SCHEMA,
 })
 
 def setup_platform(hass, config, add_entities, discovery_info=None):
@@ -36,6 +42,7 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
         'query': str(config.get(CONF_PROMETHEUS_QUERY)),
         'name': str(config.get(CONF_NAME)),
         'unit': str(config.get(CONF_UNIT_OF_MEASUREMENT)),
+        'state_class': str(config.get(CONF_STATE_CLASS)),
     }
     add_entities([PrometheusQuery(prom_data)], True)
 
@@ -49,6 +56,7 @@ class PrometheusQuery(Entity):
         self._name = prom_data["name"]
         self._state = None
         self._unit_of_measurement = prom_data["unit"]
+        self._state_class = prom_data["state_class"]
 
     @property
     def name(self):
@@ -64,6 +72,11 @@ class PrometheusQuery(Entity):
     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.