From 6590e464af2b8598df034d5e0f948ff7fa7c833b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20Hjelseth=20H=C3=B8yer?= Date: Thu, 29 Jul 2021 20:05:53 +0200 Subject: [PATCH] Integration. Add device class, last_reset, state_class (#53698) Co-authored-by: Franck Nijhof --- .../components/integration/sensor.py | 29 ++++- tests/components/integration/test_sensor.py | 116 ++++++++++++++++-- 2 files changed, 132 insertions(+), 13 deletions(-) diff --git a/homeassistant/components/integration/sensor.py b/homeassistant/components/integration/sensor.py index 0ab4ac0d2c4..dea8970f4f7 100644 --- a/homeassistant/components/integration/sensor.py +++ b/homeassistant/components/integration/sensor.py @@ -4,8 +4,16 @@ import logging import voluptuous as vol -from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity +from homeassistant.components.sensor import ( + ATTR_LAST_RESET, + DEVICE_CLASS_ENERGY, + DEVICE_CLASS_POWER, + PLATFORM_SCHEMA, + STATE_CLASS_MEASUREMENT, + SensorEntity, +) from homeassistant.const import ( + ATTR_DEVICE_CLASS, ATTR_UNIT_OF_MEASUREMENT, CONF_METHOD, CONF_NAME, @@ -20,6 +28,7 @@ from homeassistant.core import callback import homeassistant.helpers.config_validation as cv from homeassistant.helpers.event import async_track_state_change_event from homeassistant.helpers.restore_state import RestoreEntity +from homeassistant.util import dt as dt_util # mypy: allow-untyped-defs, no-check-untyped-defs @@ -115,16 +124,26 @@ class IntegrationSensor(RestoreEntity, SensorEntity): self._unit_prefix = UNIT_PREFIXES[unit_prefix] self._unit_time = UNIT_TIME[unit_time] + self._attr_state_class = STATE_CLASS_MEASUREMENT async def async_added_to_hass(self): """Handle entity which will be added.""" await super().async_added_to_hass() state = await self.async_get_last_state() + self._attr_last_reset = dt_util.utcnow() if state: try: self._state = Decimal(state.state) - except ValueError as err: + except (DecimalException, ValueError) as err: _LOGGER.warning("Could not restore last state: %s", err) + else: + last_reset = dt_util.parse_datetime( + state.attributes.get(ATTR_LAST_RESET, "") + ) + self._attr_last_reset = ( + last_reset if last_reset else dt_util.utc_from_timestamp(0) + ) + self._attr_device_class = state.attributes.get(ATTR_DEVICE_CLASS) @callback def calc_integration(event): @@ -143,7 +162,11 @@ class IntegrationSensor(RestoreEntity, SensorEntity): self._unit_of_measurement = self._unit_template.format( "" if unit is None else unit ) - + if ( + self.device_class is None + and new_state.attributes.get(ATTR_DEVICE_CLASS) == DEVICE_CLASS_POWER + ): + self._attr_device_class = DEVICE_CLASS_ENERGY try: # integration as the Riemann integral of previous measures. area = 0 diff --git a/tests/components/integration/test_sensor.py b/tests/components/integration/test_sensor.py index 3afa5c14c22..dd6bf980d0f 100644 --- a/tests/components/integration/test_sensor.py +++ b/tests/components/integration/test_sensor.py @@ -2,12 +2,22 @@ from datetime import timedelta from unittest.mock import patch -from homeassistant.const import ENERGY_KILO_WATT_HOUR, POWER_WATT, TIME_SECONDS +from homeassistant.components.sensor import STATE_CLASS_MEASUREMENT +from homeassistant.const import ( + DEVICE_CLASS_ENERGY, + DEVICE_CLASS_POWER, + ENERGY_KILO_WATT_HOUR, + POWER_WATT, + TIME_SECONDS, +) +from homeassistant.core import HomeAssistant, State from homeassistant.setup import async_setup_component import homeassistant.util.dt as dt_util +from tests.common import mock_restore_cache -async def test_state(hass): + +async def test_state(hass) -> None: """Test integration sensor state.""" config = { "sensor": { @@ -19,15 +29,25 @@ async def test_state(hass): } } - assert await async_setup_component(hass, "sensor", config) - - entity_id = config["sensor"]["source"] - hass.states.async_set(entity_id, 1, {}) - await hass.async_block_till_done() - - now = dt_util.utcnow() + timedelta(seconds=3600) + now = dt_util.utcnow() with patch("homeassistant.util.dt.utcnow", return_value=now): - hass.states.async_set(entity_id, 1, {}, force_update=True) + assert await async_setup_component(hass, "sensor", config) + + entity_id = config["sensor"]["source"] + hass.states.async_set(entity_id, 1, {}) + await hass.async_block_till_done() + + state = hass.states.get("sensor.integration") + assert state is not None + assert state.attributes.get("last_reset") == now.isoformat() + assert state.attributes.get("state_class") == STATE_CLASS_MEASUREMENT + assert "device_class" not in state.attributes + + future_now = dt_util.utcnow() + timedelta(seconds=3600) + with patch("homeassistant.util.dt.utcnow", return_value=future_now): + hass.states.async_set( + entity_id, 1, {"device_class": DEVICE_CLASS_POWER}, force_update=True + ) await hass.async_block_till_done() state = hass.states.get("sensor.integration") @@ -37,6 +57,82 @@ async def test_state(hass): assert round(float(state.state), config["sensor"]["round"]) == 1.0 assert state.attributes.get("unit_of_measurement") == ENERGY_KILO_WATT_HOUR + assert state.attributes.get("device_class") == DEVICE_CLASS_ENERGY + assert state.attributes.get("state_class") == STATE_CLASS_MEASUREMENT + assert state.attributes.get("last_reset") == now.isoformat() + + +async def test_restore_state(hass: HomeAssistant) -> None: + """Test integration sensor state is restored correctly.""" + mock_restore_cache( + hass, + ( + State( + "sensor.integration", + "100.0", + { + "last_reset": "2019-10-06T21:00:00", + "device_class": DEVICE_CLASS_ENERGY, + }, + ), + ), + ) + + config = { + "sensor": { + "platform": "integration", + "name": "integration", + "source": "sensor.power", + "unit": ENERGY_KILO_WATT_HOUR, + "round": 2, + } + } + + assert await async_setup_component(hass, "sensor", config) + await hass.async_block_till_done() + + state = hass.states.get("sensor.integration") + assert state + assert state.state == "100.00" + assert state.attributes.get("unit_of_measurement") == ENERGY_KILO_WATT_HOUR + assert state.attributes.get("device_class") == DEVICE_CLASS_ENERGY + assert state.attributes.get("last_reset") == "2019-10-06T21:00:00" + + +async def test_restore_state_failed(hass: HomeAssistant) -> None: + """Test integration sensor state is restored correctly.""" + mock_restore_cache( + hass, + ( + State( + "sensor.integration", + "INVALID", + { + "last_reset": "2019-10-06T21:00:00.000000", + }, + ), + ), + ) + + config = { + "sensor": { + "platform": "integration", + "name": "integration", + "source": "sensor.power", + "unit": ENERGY_KILO_WATT_HOUR, + } + } + + assert await async_setup_component(hass, "sensor", config) + await hass.async_block_till_done() + + state = hass.states.get("sensor.integration") + assert state + assert state.state == "0" + assert state.attributes.get("unit_of_measurement") == ENERGY_KILO_WATT_HOUR + assert state.attributes.get("state_class") == STATE_CLASS_MEASUREMENT + assert state.attributes.get("last_reset") != "2019-10-06T21:00:00" + assert "device_class" not in state.attributes async def test_trapezoidal(hass):