From 1692fab66477e26aa93264e4a6b0e29f32d123fc Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Tue, 14 Dec 2021 18:32:33 +0100 Subject: [PATCH] Use new enums in integration (#61803) Co-authored-by: epenet --- homeassistant/components/integration/sensor.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/homeassistant/components/integration/sensor.py b/homeassistant/components/integration/sensor.py index 463cb3b4e05..ee7b4b8e86e 100644 --- a/homeassistant/components/integration/sensor.py +++ b/homeassistant/components/integration/sensor.py @@ -5,11 +5,10 @@ import logging import voluptuous as vol from homeassistant.components.sensor import ( - DEVICE_CLASS_ENERGY, - DEVICE_CLASS_POWER, PLATFORM_SCHEMA, - STATE_CLASS_TOTAL, + SensorDeviceClass, SensorEntity, + SensorStateClass, ) from homeassistant.const import ( ATTR_DEVICE_CLASS, @@ -119,7 +118,7 @@ class IntegrationSensor(RestoreEntity, SensorEntity): self._unit_of_measurement = unit_of_measurement self._unit_prefix = UNIT_PREFIXES[unit_prefix] self._unit_time = UNIT_TIME[unit_time] - self._attr_state_class = STATE_CLASS_TOTAL + self._attr_state_class = SensorStateClass.TOTAL async def async_added_to_hass(self): """Handle entity which will be added.""" @@ -149,9 +148,10 @@ class IntegrationSensor(RestoreEntity, SensorEntity): ) if ( self.device_class is None - and new_state.attributes.get(ATTR_DEVICE_CLASS) == DEVICE_CLASS_POWER + and new_state.attributes.get(ATTR_DEVICE_CLASS) + == SensorDeviceClass.POWER ): - self._attr_device_class = DEVICE_CLASS_ENERGY + self._attr_device_class = SensorDeviceClass.ENERGY if ( old_state is None