Use new enums in energy (#61386)

Co-authored-by: epenet <epenet@users.noreply.github.com>

homeassistant/components/energy/sensor.py

@@ -10,11 +10,9 @@ from typing import Any, Final, Literal, TypeVar, cast
 from homeassistant.components.sensor import (
     ATTR_LAST_RESET,
     ATTR_STATE_CLASS,
-    DEVICE_CLASS_MONETARY,
-    STATE_CLASS_MEASUREMENT,
-    STATE_CLASS_TOTAL,
-    STATE_CLASS_TOTAL_INCREASING,
+    SensorDeviceClass,
     SensorEntity,
+    SensorStateClass,
 )
 from homeassistant.components.sensor.recorder import reset_detected
 from homeassistant.const import (
@@ -22,7 +20,6 @@ from homeassistant.const import (
     ENERGY_KILO_WATT_HOUR,
     ENERGY_MEGA_WATT_HOUR,
     ENERGY_WATT_HOUR,
-    ENTITY_CATEGORY_SYSTEM,
     VOLUME_CUBIC_METERS,
 )
 from homeassistant.core import (
@@ -32,6 +29,7 @@ from homeassistant.core import (
     split_entity_id,
     valid_entity_id,
 )
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_track_state_change_event
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
@@ -41,9 +39,9 @@ from .const import DOMAIN
 from .data import EnergyManager, async_get_manager
 
 SUPPORTED_STATE_CLASSES = [
-    STATE_CLASS_MEASUREMENT,
-    STATE_CLASS_TOTAL,
-    STATE_CLASS_TOTAL_INCREASING,
+    SensorStateClass.MEASUREMENT,
+    SensorStateClass.TOTAL,
+    SensorStateClass.TOTAL_INCREASING,
 ]
 VALID_ENERGY_UNITS = [ENERGY_WATT_HOUR, ENERGY_KILO_WATT_HOUR, ENERGY_MEGA_WATT_HOUR]
 VALID_ENERGY_UNITS_GAS = [VOLUME_CUBIC_METERS] + VALID_ENERGY_UNITS
@@ -215,7 +213,7 @@ class EnergyCostSensor(SensorEntity):
     utility.
     """
 
-    _attr_entity_category = ENTITY_CATEGORY_SYSTEM
+    _attr_entity_category = EntityCategory.SYSTEM
     _wrong_state_class_reported = False
     _wrong_unit_reported = False
 
@@ -231,8 +229,8 @@ class EnergyCostSensor(SensorEntity):
         self.entity_id = (
             f"{config[adapter.entity_energy_key]}_{adapter.entity_id_suffix}"
         )
-        self._attr_device_class = DEVICE_CLASS_MONETARY
-        self._attr_state_class = STATE_CLASS_TOTAL
+        self._attr_device_class = SensorDeviceClass.MONETARY
+        self._attr_state_class = SensorStateClass.TOTAL
         self._config = config
         self._last_energy_sensor_state: State | None = None
         # add_finished is set when either of async_added_to_hass or add_to_platform_abort
@@ -267,9 +265,9 @@ class EnergyCostSensor(SensorEntity):
                 )
             return
 
-        # last_reset must be set if the sensor is STATE_CLASS_MEASUREMENT
+        # last_reset must be set if the sensor is SensorStateClass.MEASUREMENT
         if (
-            state_class == STATE_CLASS_MEASUREMENT
+            state_class == SensorStateClass.MEASUREMENT
             and ATTR_LAST_RESET not in energy_state.attributes
         ):
             return
@@ -337,14 +335,16 @@ class EnergyCostSensor(SensorEntity):
                 )
             return
 
-        if state_class != STATE_CLASS_TOTAL_INCREASING and energy_state.attributes.get(
-            ATTR_LAST_RESET
-        ) != self._last_energy_sensor_state.attributes.get(ATTR_LAST_RESET):
+        if (
+            state_class != SensorStateClass.TOTAL_INCREASING
+            and energy_state.attributes.get(ATTR_LAST_RESET)
+            != self._last_energy_sensor_state.attributes.get(ATTR_LAST_RESET)
+        ):
             # Energy meter was reset, reset cost sensor too
             energy_state_copy = copy.copy(energy_state)
             energy_state_copy.state = "0.0"
             self._reset(energy_state_copy)
-        elif state_class == STATE_CLASS_TOTAL_INCREASING and reset_detected(
+        elif state_class == SensorStateClass.TOTAL_INCREASING and reset_detected(
             self.hass,
             cast(str, self._config[self._adapter.entity_energy_key]),
             energy,

homeassistant/components/energy/validate.py

@@ -21,19 +21,22 @@ from homeassistant.core import HomeAssistant, callback, valid_entity_id
 from . import data
 from .const import DOMAIN
 
-ENERGY_USAGE_DEVICE_CLASSES = (sensor.DEVICE_CLASS_ENERGY,)
+ENERGY_USAGE_DEVICE_CLASSES = (sensor.SensorDeviceClass.ENERGY,)
 ENERGY_USAGE_UNITS = {
-    sensor.DEVICE_CLASS_ENERGY: (ENERGY_KILO_WATT_HOUR, ENERGY_WATT_HOUR)
+    sensor.SensorDeviceClass.ENERGY: (ENERGY_KILO_WATT_HOUR, ENERGY_WATT_HOUR)
 }
 ENERGY_PRICE_UNITS = tuple(
     f"/{unit}" for units in ENERGY_USAGE_UNITS.values() for unit in units
 )
 ENERGY_UNIT_ERROR = "entity_unexpected_unit_energy"
 ENERGY_PRICE_UNIT_ERROR = "entity_unexpected_unit_energy_price"
-GAS_USAGE_DEVICE_CLASSES = (sensor.DEVICE_CLASS_ENERGY, sensor.DEVICE_CLASS_GAS)
+GAS_USAGE_DEVICE_CLASSES = (
+    sensor.SensorDeviceClass.ENERGY,
+    sensor.SensorDeviceClass.GAS,
+)
 GAS_USAGE_UNITS = {
-    sensor.DEVICE_CLASS_ENERGY: (ENERGY_WATT_HOUR, ENERGY_KILO_WATT_HOUR),
-    sensor.DEVICE_CLASS_GAS: (VOLUME_CUBIC_METERS, VOLUME_CUBIC_FEET),
+    sensor.SensorDeviceClass.ENERGY: (ENERGY_WATT_HOUR, ENERGY_KILO_WATT_HOUR),
+    sensor.SensorDeviceClass.GAS: (VOLUME_CUBIC_METERS, VOLUME_CUBIC_FEET),
 }
 GAS_PRICE_UNITS = tuple(
     f"/{unit}" for units in GAS_USAGE_UNITS.values() for unit in units
@@ -141,9 +144,9 @@ def _async_validate_usage_stat(
     state_class = state.attributes.get(sensor.ATTR_STATE_CLASS)
 
     allowed_state_classes = [
-        sensor.STATE_CLASS_MEASUREMENT,
-        sensor.STATE_CLASS_TOTAL,
-        sensor.STATE_CLASS_TOTAL_INCREASING,
+        sensor.SensorStateClass.MEASUREMENT,
+        sensor.SensorStateClass.TOTAL,
+        sensor.SensorStateClass.TOTAL_INCREASING,
     ]
     if state_class not in allowed_state_classes:
         result.append(
@@ -155,7 +158,7 @@ def _async_validate_usage_stat(
         )
 
     if (
-        state_class == sensor.STATE_CLASS_MEASUREMENT
+        state_class == sensor.SensorStateClass.MEASUREMENT
         and sensor.ATTR_LAST_RESET not in state.attributes
     ):
         result.append(
@@ -221,9 +224,9 @@ def _async_validate_cost_stat(
     state_class = state.attributes.get("state_class")
 
     supported_state_classes = [
-        sensor.STATE_CLASS_MEASUREMENT,
-        sensor.STATE_CLASS_TOTAL,
-        sensor.STATE_CLASS_TOTAL_INCREASING,
+        sensor.SensorStateClass.MEASUREMENT,
+        sensor.SensorStateClass.TOTAL,
+        sensor.SensorStateClass.TOTAL_INCREASING,
     ]
     if state_class not in supported_state_classes:
         result.append(
@@ -231,7 +234,7 @@ def _async_validate_cost_stat(
         )
 
     if (
-        state_class == sensor.STATE_CLASS_MEASUREMENT
+        state_class == sensor.SensorStateClass.MEASUREMENT
         and sensor.ATTR_LAST_RESET not in state.attributes
     ):
         result.append(