diff --git a/homeassistant/components/srp_energy/sensor.py b/homeassistant/components/srp_energy/sensor.py index 97b65840e83..4a5e3c33748 100644 --- a/homeassistant/components/srp_energy/sensor.py +++ b/homeassistant/components/srp_energy/sensor.py @@ -5,8 +5,12 @@ import logging import async_timeout from requests.exceptions import ConnectionError as ConnectError, HTTPError, Timeout -from homeassistant.components.sensor import SensorEntity -from homeassistant.const import ATTR_ATTRIBUTION, ENERGY_KILO_WATT_HOUR +from homeassistant.components.sensor import STATE_CLASS_TOTAL_INCREASING, SensorEntity +from homeassistant.const import ( + ATTR_ATTRIBUTION, + DEVICE_CLASS_ENERGY, + ENERGY_KILO_WATT_HOUR, +) from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from .const import ( @@ -137,6 +141,16 @@ class SrpEntity(SensorEntity): """Return if entity is available.""" return self.coordinator.last_update_success + @property + def device_class(self): + """Return the device class.""" + return DEVICE_CLASS_ENERGY + + @property + def state_class(self): + """Return the state class.""" + return STATE_CLASS_TOTAL_INCREASING + async def async_added_to_hass(self): """When entity is added to hass.""" self.async_on_remove( diff --git a/tests/components/srp_energy/test_sensor.py b/tests/components/srp_energy/test_sensor.py index 3e830b7fc93..5a4585cc8e5 100644 --- a/tests/components/srp_energy/test_sensor.py +++ b/tests/components/srp_energy/test_sensor.py @@ -1,6 +1,7 @@ """Tests for the srp_energy sensor platform.""" from unittest.mock import MagicMock +from homeassistant.components.sensor import STATE_CLASS_TOTAL_INCREASING from homeassistant.components.srp_energy.const import ( ATTRIBUTION, DEFAULT_NAME, @@ -10,7 +11,11 @@ from homeassistant.components.srp_energy.const import ( SRP_ENERGY_DOMAIN, ) from homeassistant.components.srp_energy.sensor import SrpEntity, async_setup_entry -from homeassistant.const import ATTR_ATTRIBUTION, ENERGY_KILO_WATT_HOUR +from homeassistant.const import ( + ATTR_ATTRIBUTION, + DEVICE_CLASS_ENERGY, + ENERGY_KILO_WATT_HOUR, +) async def test_async_setup_entry(hass): @@ -94,6 +99,8 @@ async def test_srp_entity(hass): assert srp_entity.should_poll is False assert srp_entity.extra_state_attributes[ATTR_ATTRIBUTION] == ATTRIBUTION assert srp_entity.available is not None + assert srp_entity.device_class == DEVICE_CLASS_ENERGY + assert srp_entity.state_class == STATE_CLASS_TOTAL_INCREASING await srp_entity.async_added_to_hass() assert srp_entity.state is not None