Integration. Add device class, last_reset, state_class (#53698)

Co-authored-by: Franck Nijhof <git@frenck.dev>

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

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)
+    now = dt_util.utcnow()
-
-    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)
     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):