From 1ae200fe89280d14d02d3ab1c84a99c459d8974c Mon Sep 17 00:00:00 2001
From: Diogo Gomes <dgomes@ua.pt>
Date: Sun, 25 Feb 2024 16:04:42 +0000
Subject: [PATCH] address 2024.1 change from magic numbers to enum

---
 custom_components/erse/entity.py | 4 ++--
 custom_components/erse/sensor.py | 9 ++++-----
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/custom_components/erse/entity.py b/custom_components/erse/entity.py
index 2e2ef40..160bfb2 100644
--- a/custom_components/erse/entity.py
+++ b/custom_components/erse/entity.py
@@ -1,7 +1,7 @@
 """Support for ERSE entities."""
 from __future__ import annotations
 
-from homeassistant.components.sensor import STATE_CLASS_TOTAL, SensorDeviceClass
+from homeassistant.components.sensor import SensorStateClass, SensorDeviceClass
 from homeassistant.const import CURRENCY_EURO
 from homeassistant.helpers.entity import Entity, DeviceInfo
 from pyerse.comercializador import Comercializador
@@ -38,6 +38,6 @@ class ERSEMoneyEntity(ERSEEntity):
     """Defines a monetary ERSE entity."""
 
     _attr_device_class = SensorDeviceClass.MONETARY
-    _attr_state_class = STATE_CLASS_TOTAL
+    _attr_state_class = SensorStateClass.TOTAL
     _attr_native_unit_of_measurement = CURRENCY_EURO
     _attr_suggested_display_precision = COST_PRECISION
diff --git a/custom_components/erse/sensor.py b/custom_components/erse/sensor.py
index d8185ee..440683b 100644
--- a/custom_components/erse/sensor.py
+++ b/custom_components/erse/sensor.py
@@ -18,11 +18,10 @@
 from homeassistant.const import (
     ATTR_ENTITY_ID,
     ATTR_UNIT_OF_MEASUREMENT,
-    ENERGY_KILO_WATT_HOUR,
-    ENERGY_WATT_HOUR,
     EVENT_HOMEASSISTANT_START,
     STATE_UNAVAILABLE,
     STATE_UNKNOWN,
+    UnitOfEnergy
 )
 from homeassistant.core import callback
 from homeassistant.helpers.entity import Entity
@@ -222,7 +221,7 @@ def __init__(self, hass, entry_id, export_entity, tariff, meter_entities):
 
         self._attr_name = f"{tariff.value} Net"
         self._attr_unique_id = slugify(f"{entry_id} {tariff} netmeter")
-        self._attr_native_unit_of_measurement = ENERGY_KILO_WATT_HOUR
+        self._attr_native_unit_of_measurement = UnitOfEnergy.KILO_WATT_HOUR
 
         self._export_entity = export_entity
         self._tariff = tariff
@@ -436,8 +435,8 @@ async def calc_costs(meter_state):
                 and ATTR_UNIT_OF_MEASUREMENT in meter_state.attributes
                 and meter_state.attributes[ATTR_UNIT_OF_MEASUREMENT]
                 in [
-                    ENERGY_WATT_HOUR,
-                    ENERGY_KILO_WATT_HOUR,
+                    UnitOfEnergy.WATT_HOUR,
+                    UnitOfEnergy.KILO_WATT_HOUR,
                 ]
             ):
                 if meter_state.state in [STATE_UNAVAILABLE, STATE_UNKNOWN]: