Calculate device class as soon as it is known in integral (#119940)

tests/components/integration/test_sensor.py

@@ -5,10 +5,12 @@ from typing import Any
 
 from freezegun import freeze_time
 import pytest
+from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.components.integration.const import DOMAIN
 from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
 from homeassistant.const import (
+    ATTR_DEVICE_CLASS,
     ATTR_UNIT_OF_MEASUREMENT,
     STATE_UNAVAILABLE,
     STATE_UNKNOWN,
@@ -17,6 +19,7 @@ from homeassistant.const import (
     UnitOfInformation,
     UnitOfPower,
     UnitOfTime,
+    UnitOfVolumeFlowRate,
 )
 from homeassistant.core import HomeAssistant, State
 from homeassistant.helpers import (
@@ -36,6 +39,52 @@ from tests.common import (
 DEFAULT_MAX_SUB_INTERVAL = {"minutes": 1}
 
 
+@pytest.mark.parametrize(
+    ("unit_of_measurement", "device_class", "unit_time"),
+    [
+        (UnitOfPower.KILO_WATT, SensorDeviceClass.POWER, "h"),
+        (UnitOfPower.KILO_WATT, None, "h"),
+        (UnitOfPower.BTU_PER_HOUR, SensorDeviceClass.POWER, "h"),
+        (
+            UnitOfVolumeFlowRate.CUBIC_FEET_PER_MINUTE,
+            SensorDeviceClass.VOLUME_FLOW_RATE,
+            "min",
+        ),
+    ],
+)
+async def test_initial_state(
+    hass: HomeAssistant,
+    unit_of_measurement: str,
+    device_class: SensorDeviceClass,
+    unit_time: str,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test integration sensor state."""
+    config = {
+        "sensor": {
+            "platform": "integration",
+            "name": "integration",
+            "source": "sensor.source",
+            "round": 2,
+            "method": "left",
+            "unit_time": unit_time,
+        }
+    }
+
+    assert await async_setup_component(hass, "sensor", config)
+    hass.states.async_set(
+        "sensor.source",
+        "1",
+        {
+            ATTR_DEVICE_CLASS: device_class,
+            ATTR_UNIT_OF_MEASUREMENT: unit_of_measurement,
+        },
+    )
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.integration") == snapshot
+
+
 @pytest.mark.parametrize("method", ["trapezoidal", "left", "right"])
 async def test_state(hass: HomeAssistant, method) -> None:
     """Test integration sensor state."""
@@ -49,13 +98,23 @@ async def test_state(hass: HomeAssistant, method) -> None:
         }
     }
 
+    assert await async_setup_component(hass, "sensor", config)
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.integration")
+    assert state is not None
+    assert state.attributes.get("state_class") is SensorStateClass.TOTAL
+    assert "device_class" not in state.attributes
+
     now = dt_util.utcnow()
     with freeze_time(now):
-        assert await async_setup_component(hass, "sensor", config)
-
         entity_id = config["sensor"]["source"]
         hass.states.async_set(
-            entity_id, 1, {ATTR_UNIT_OF_MEASUREMENT: UnitOfPower.KILO_WATT}
+            entity_id,
+            1,
+            {
+                ATTR_UNIT_OF_MEASUREMENT: UnitOfPower.KILO_WATT,
+            },
         )
         await hass.async_block_till_done()