From 2582f6f5a9ee2a3e146bae07eae9d0208cbed10a Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:42:59 +0200 Subject: [PATCH] Use EntityStateAttribute enum in integration (#175933) Co-authored-by: Claude Opus 4.8 (1M context) --- .../components/integration/config_flow.py | 14 ++++++-------- homeassistant/components/integration/sensor.py | 10 ++++++---- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/homeassistant/components/integration/config_flow.py b/homeassistant/components/integration/config_flow.py index 7d00d2161cc0..718fee7893c5 100644 --- a/homeassistant/components/integration/config_flow.py +++ b/homeassistant/components/integration/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_METHOD, - CONF_NAME, - UnitOfTime, -) +from homeassistant.const import CONF_METHOD, CONF_NAME, EntityStateAttribute, UnitOfTime from homeassistant.core import callback from homeassistant.helpers import selector from homeassistant.helpers.schema_config_entry_flow import ( @@ -62,13 +57,16 @@ def entity_selector_compatible( """Return an entity selector which compatible entities.""" current = handler.hass.states.get(handler.options[CONF_SOURCE_SENSOR]) 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 ] diff --git a/homeassistant/components/integration/sensor.py b/homeassistant/components/integration/sensor.py index 83b93c4191d0..edd7c5d0c1d0 100644 --- a/homeassistant/components/integration/sensor.py +++ b/homeassistant/components/integration/sensor.py @@ -20,12 +20,11 @@ from homeassistant.components.sensor import ( ) from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( - ATTR_DEVICE_CLASS, - ATTR_UNIT_OF_MEASUREMENT, CONF_METHOD, CONF_NAME, CONF_UNIQUE_ID, STATE_UNAVAILABLE, + EntityStateAttribute, UnitOfTime, ) from homeassistant.core import ( @@ -389,7 +388,9 @@ class IntegrationSensor(RestoreSensor): return device_class def _derive_and_set_attributes_from_state(self, source_state: State) -> None: - source_unit = source_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) + source_unit = source_state.attributes.get( + EntityStateAttribute.UNIT_OF_MEASUREMENT + ) if source_unit is not None: self._unit_of_measurement = self._calculate_unit(source_unit) else: @@ -397,7 +398,8 @@ class IntegrationSensor(RestoreSensor): self._unit_of_measurement = None self._attr_device_class = self._calculate_device_class( - source_state.attributes.get(ATTR_DEVICE_CLASS), self.unit_of_measurement + source_state.attributes.get(EntityStateAttribute.DEVICE_CLASS), + self.unit_of_measurement, ) if self._attr_device_class: # Remove this sensors icon default and allow