address 2024.1 change from magic numbers to enum

dgomes · Feb 25, 2024 · 1ae200f · 1ae200f
1 parent 5759d16
commit 1ae200f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git 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
@@ -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]: