Improve integration sensor's time unit handling (#72759)

homeassistant/components/integration/sensor.py

@@ -154,17 +154,26 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
         self._method = integration_method
 
         self._attr_name = name if name is not None else f"{source_entity} integral"
-        self._unit_template = (
-            f"{'' if unit_prefix is None else unit_prefix}{{}}{unit_time}"
-        )
+        self._unit_template = f"{'' if unit_prefix is None else unit_prefix}{{}}"
         self._unit_of_measurement = None
         self._unit_prefix = UNIT_PREFIXES[unit_prefix]
         self._unit_time = UNIT_TIME[unit_time]
+        self._unit_time_str = unit_time
         self._attr_state_class = SensorStateClass.TOTAL
         self._attr_icon = "mdi:chart-histogram"
         self._attr_should_poll = False
         self._attr_extra_state_attributes = {ATTR_SOURCE_ID: source_entity}
 
+    def _unit(self, source_unit: str) -> str:
+        """Derive unit from the source sensor, SI prefix and time unit."""
+        unit_time = self._unit_time_str
+        if source_unit.endswith(f"/{unit_time}"):
+            integral_unit = source_unit[0 : (-(1 + len(unit_time)))]
+        else:
+            integral_unit = f"{source_unit}{unit_time}"
+
+        return self._unit_template.format(integral_unit)
+
     async def async_added_to_hass(self):
         """Handle entity which will be added."""
         await super().async_added_to_hass()
@@ -203,7 +212,7 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
             update_state = False
             unit = new_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
             if unit is not None:
-                new_unit_of_measurement = self._unit_template.format(unit)
+                new_unit_of_measurement = self._unit(unit)
                 if self._unit_of_measurement != new_unit_of_measurement:
                     self._unit_of_measurement = new_unit_of_measurement
                     update_state = True

tests/components/integration/test_sensor.py

@@ -5,12 +5,15 @@ from unittest.mock import patch
 from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
 from homeassistant.const import (
     ATTR_UNIT_OF_MEASUREMENT,
+    DATA_KILOBYTES,
+    DATA_RATE_BYTES_PER_SECOND,
     ENERGY_KILO_WATT_HOUR,
     ENERGY_WATT_HOUR,
     POWER_KILO_WATT,
     POWER_WATT,
     STATE_UNAVAILABLE,
     STATE_UNKNOWN,
+    TIME_HOURS,
     TIME_SECONDS,
 )
 from homeassistant.core import HomeAssistant, State
@@ -300,7 +303,9 @@ async def test_suffix(hass):
     assert await async_setup_component(hass, "sensor", config)
 
     entity_id = config["sensor"]["source"]
-    hass.states.async_set(entity_id, 1000, {ATTR_UNIT_OF_MEASUREMENT: POWER_KILO_WATT})
+    hass.states.async_set(
+        entity_id, 1000, {ATTR_UNIT_OF_MEASUREMENT: DATA_RATE_BYTES_PER_SECOND}
+    )
     await hass.async_block_till_done()
 
     now = dt_util.utcnow() + timedelta(seconds=10)
@@ -308,7 +313,7 @@ async def test_suffix(hass):
         hass.states.async_set(
             entity_id,
             1000,
-            {ATTR_UNIT_OF_MEASUREMENT: POWER_KILO_WATT},
+            {ATTR_UNIT_OF_MEASUREMENT: DATA_RATE_BYTES_PER_SECOND},
             force_update=True,
         )
         await hass.async_block_till_done()
@@ -318,6 +323,43 @@ async def test_suffix(hass):
 
     # Testing a network speed sensor at 1000 bytes/s over 10s  = 10kbytes
     assert round(float(state.state)) == 10
+    assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == DATA_KILOBYTES
+
+
+async def test_suffix_2(hass):
+    """Test integration sensor state."""
+    config = {
+        "sensor": {
+            "platform": "integration",
+            "name": "integration",
+            "source": "sensor.cubic_meters_per_hour",
+            "round": 2,
+            "unit_time": TIME_HOURS,
+        }
+    }
+
+    assert await async_setup_component(hass, "sensor", config)
+
+    entity_id = config["sensor"]["source"]
+    hass.states.async_set(entity_id, 1000, {ATTR_UNIT_OF_MEASUREMENT: "m³/h"})
+    await hass.async_block_till_done()
+
+    now = dt_util.utcnow() + timedelta(hours=1)
+    with patch("homeassistant.util.dt.utcnow", return_value=now):
+        hass.states.async_set(
+            entity_id,
+            1000,
+            {ATTR_UNIT_OF_MEASUREMENT: "m³/h"},
+            force_update=True,
+        )
+        await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.integration")
+    assert state is not None
+
+    # Testing a flow sensor at 1000 m³/h over 1h = 1000 m³
+    assert round(float(state.state)) == 1000
+    assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == "m³"
 
 
 async def test_units(hass):