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

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

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):