Use state attribute enums in derivative (#175848)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

homeassistant/components/derivative/config_flow.py

@@ -8,12 +8,7 @@ import voluptuous as vol
 from homeassistant.components.counter import DOMAIN as COUNTER_DOMAIN
 from homeassistant.components.input_number import DOMAIN as INPUT_NUMBER_DOMAIN
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
-from homeassistant.const import (
-    ATTR_UNIT_OF_MEASUREMENT,
-    CONF_NAME,
-    CONF_SOURCE,
-    UnitOfTime,
-)
+from homeassistant.const import CONF_NAME, CONF_SOURCE, EntityStateAttribute, UnitOfTime
 from homeassistant.core import callback
 from homeassistant.helpers import selector
 from homeassistant.helpers.schema_config_entry_flow import (
@@ -59,13 +54,16 @@ def entity_selector_compatible(
     """Return an entity selector which compatible entities."""
     current = handler.hass.states.get(handler.options[CONF_SOURCE])
     unit_of_measurement = (
-        current.attributes.get(ATTR_UNIT_OF_MEASUREMENT) if current else None
+        current.attributes.get(EntityStateAttribute.UNIT_OF_MEASUREMENT)
+        if current
+        else None
     )
 
     entities = [
         ent.entity_id
         for ent in handler.hass.states.async_all(ALLOWED_DOMAINS)
-        if ent.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == unit_of_measurement
+        if ent.attributes.get(EntityStateAttribute.UNIT_OF_MEASUREMENT)
+        == unit_of_measurement
         and ent.domain in ALLOWED_DOMAINS
     ]
 

homeassistant/components/derivative/sensor.py

@@ -8,23 +8,22 @@ from typing import override
 import voluptuous as vol
 
 from homeassistant.components.sensor import (
-    ATTR_STATE_CLASS,
     DEVICE_CLASS_UNITS,
     PLATFORM_SCHEMA as SENSOR_PLATFORM_SCHEMA,
     RestoreSensor,
     SensorDeviceClass,
     SensorEntity,
+    SensorEntityCapabilityAttribute,
     SensorStateClass,
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
-    ATTR_DEVICE_CLASS,
-    ATTR_UNIT_OF_MEASUREMENT,
     CONF_NAME,
     CONF_SOURCE,
     CONF_UNIQUE_ID,
     STATE_UNAVAILABLE,
     STATE_UNKNOWN,
+    EntityStateAttribute,
     Platform,
     UnitOfTime,
 )
@@ -243,7 +242,9 @@ class DerivativeSensor(RestoreSensor, SensorEntity):
         if not source_state:
             return
 
-        source_class_raw = source_state.attributes.get(ATTR_DEVICE_CLASS)
+        source_class_raw = source_state.attributes.get(
+            EntityStateAttribute.DEVICE_CLASS
+        )
         source_class: SensorDeviceClass | None = None
         if isinstance(source_class_raw, str):
             try:
@@ -252,7 +253,9 @@ class DerivativeSensor(RestoreSensor, SensorEntity):
                 source_class = None
         if self._string_unit_prefix is not None and self._string_unit_time is not None:
             original_unit = self._attr_native_unit_of_measurement
-            source_unit = source_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
+            source_unit = source_state.attributes.get(
+                EntityStateAttribute.UNIT_OF_MEASUREMENT
+            )
             if (
                 (
                     source_class
@@ -366,7 +369,9 @@ class DerivativeSensor(RestoreSensor, SensorEntity):
 
         last_state = await self.async_get_last_state()
         if last_state:
-            self._attr_device_class = last_state.attributes.get(ATTR_DEVICE_CLASS)
+            self._attr_device_class = last_state.attributes.get(
+                EntityStateAttribute.DEVICE_CLASS
+            )
 
     @override
     async def async_added_to_hass(self) -> None:
@@ -540,7 +545,7 @@ class DerivativeSensor(RestoreSensor, SensorEntity):
             # A negative derivative for a total increasing sensor likely indicates the
             # sensor has been reset. To prevent inaccurate data, discard this sample.
             if (
-                new_state.attributes.get(ATTR_STATE_CLASS)
+                new_state.attributes.get(SensorEntityCapabilityAttribute.STATE_CLASS)
                 == SensorStateClass.TOTAL_INCREASING
                 and new_derivative < 0
             ):