Better handling of source sensor unavailability in Riemman Integration (#93137)

* refactor and increase coverage

* fix log order

homeassistant/components/integration/sensor.py

@@ -197,25 +197,23 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
             old_state: State | None = event.data.get("old_state")
             new_state: State | None = event.data.get("new_state")
 
-            if (
-                source_state := self.hass.states.get(self._sensor_source_id)
-            ) is None or source_state.state == STATE_UNAVAILABLE:
-                self._attr_available = False
-                self.async_write_ha_state()
-                return
-
-            self._attr_available = True
-
-            if new_state is None or new_state.state in (
-                STATE_UNKNOWN,
-                STATE_UNAVAILABLE,
-            ):
-                return
-
             # We may want to update our state before an early return,
             # based on the source sensor's unit_of_measurement
             # or device_class.
             update_state = False
+
+            if (
+                source_state := self.hass.states.get(self._sensor_source_id)
+            ) is None or source_state.state == STATE_UNAVAILABLE:
+                self._attr_available = False
+                update_state = True
+            else:
+                self._attr_available = True
+
+            if old_state is None or new_state is None:
+                # we can't calculate the elapsed time, so we can't calculate the integral
+                return
+
             unit = new_state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
             if unit is not None:
                 new_unit_of_measurement = self._unit(unit)
@@ -235,31 +233,53 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
             if update_state:
                 self.async_write_ha_state()
 
-            if old_state is None or old_state.state in (
-                STATE_UNKNOWN,
-                STATE_UNAVAILABLE,
-            ):
-                return
-
             try:
                 # integration as the Riemann integral of previous measures.
-                area = Decimal(0)
                 elapsed_time = (
                     new_state.last_updated - old_state.last_updated
                 ).total_seconds()
 
-                if self._method == METHOD_TRAPEZOIDAL:
+                if (
+                    self._method == METHOD_TRAPEZOIDAL
+                    and new_state.state
+                    not in (
+                        STATE_UNKNOWN,
+                        STATE_UNAVAILABLE,
+                    )
+                    and old_state.state
+                    not in (
+                        STATE_UNKNOWN,
+                        STATE_UNAVAILABLE,
+                    )
+                ):
                     area = (
                         (Decimal(new_state.state) + Decimal(old_state.state))
                         * Decimal(elapsed_time)
                         / 2
                     )
-                elif self._method == METHOD_LEFT:
+                elif self._method == METHOD_LEFT and old_state.state not in (
+                    STATE_UNKNOWN,
+                    STATE_UNAVAILABLE,
+                ):
                     area = Decimal(old_state.state) * Decimal(elapsed_time)
-                elif self._method == METHOD_RIGHT:
+                elif self._method == METHOD_RIGHT and new_state.state not in (
+                    STATE_UNKNOWN,
+                    STATE_UNAVAILABLE,
+                ):
                     area = Decimal(new_state.state) * Decimal(elapsed_time)
+                else:
+                    _LOGGER.debug(
+                        "Could not apply method %s to %s -> %s",
+                        self._method,
+                        old_state.state,
+                        new_state.state,
+                    )
+                    return
 
                 integral = area / (self._unit_prefix * self._unit_time)
+                _LOGGER.debug(
+                    "area = %s, integral = %s state = %s", area, integral, self._state
+                )
                 assert isinstance(integral, Decimal)
             except ValueError as err:
                 _LOGGER.warning("While calculating integration: %s", err)

tests/components/integration/test_sensor.py

@@ -2,6 +2,8 @@
 from datetime import timedelta
 from unittest.mock import patch
 
+import pytest
+
 from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
 from homeassistant.const import (
     ATTR_UNIT_OF_MEASUREMENT,
@@ -20,7 +22,8 @@ import homeassistant.util.dt as dt_util
 from tests.common import mock_restore_cache
 
 
-async def test_state(hass: HomeAssistant) -> None:
+@pytest.mark.parametrize("method", ["trapezoidal", "left", "right"])
+async def test_state(hass: HomeAssistant, method) -> None:
     """Test integration sensor state."""
     config = {
         "sensor": {
@@ -28,6 +31,7 @@ async def test_state(hass: HomeAssistant) -> None:
             "name": "integration",
             "source": "sensor.power",
             "round": 2,
+            "method": method,
         }
     }
 
@@ -46,8 +50,8 @@ async def test_state(hass: HomeAssistant) -> None:
     assert state.attributes.get("state_class") is SensorStateClass.TOTAL
     assert "device_class" not in state.attributes
 
-    future_now = dt_util.utcnow() + timedelta(seconds=3600)
+    now += timedelta(seconds=3600)
-    with patch("homeassistant.util.dt.utcnow", return_value=future_now):
+    with patch("homeassistant.util.dt.utcnow", return_value=now):
         hass.states.async_set(
             entity_id,
             1,
@@ -69,6 +73,62 @@ async def test_state(hass: HomeAssistant) -> None:
     assert state.attributes.get("device_class") == SensorDeviceClass.ENERGY
     assert state.attributes.get("state_class") is SensorStateClass.TOTAL
 
+    # 1 hour after last update, power sensor is unavailable
+    now += timedelta(seconds=3600)
+    with patch("homeassistant.util.dt.utcnow", return_value=now):
+        hass.states.async_set(
+            entity_id,
+            STATE_UNAVAILABLE,
+            {
+                "device_class": SensorDeviceClass.POWER,
+                ATTR_UNIT_OF_MEASUREMENT: UnitOfPower.KILO_WATT,
+            },
+            force_update=True,
+        )
+        await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.integration")
+    assert state.state == STATE_UNAVAILABLE
+
+    # 1 hour after last update, power sensor is back to normal at 2 KiloWatts and stays for 1 hour += 2kWh
+    now += timedelta(seconds=3600)
+    with patch("homeassistant.util.dt.utcnow", return_value=now):
+        hass.states.async_set(
+            entity_id,
+            2,
+            {
+                "device_class": SensorDeviceClass.POWER,
+                ATTR_UNIT_OF_MEASUREMENT: UnitOfPower.KILO_WATT,
+            },
+            force_update=True,
+        )
+        await hass.async_block_till_done()
+    state = hass.states.get("sensor.integration")
+    assert (
+        round(float(state.state), config["sensor"]["round"]) == 3.0
+        if method == "right"
+        else 1.0
+    )
+
+    now += timedelta(seconds=3600)
+    with patch("homeassistant.util.dt.utcnow", return_value=now):
+        hass.states.async_set(
+            entity_id,
+            2,
+            {
+                "device_class": SensorDeviceClass.POWER,
+                ATTR_UNIT_OF_MEASUREMENT: UnitOfPower.KILO_WATT,
+            },
+            force_update=True,
+        )
+        await hass.async_block_till_done()
+    state = hass.states.get("sensor.integration")
+    assert (
+        round(float(state.state), config["sensor"]["round"]) == 5.0
+        if method == "right"
+        else 3.0
+    )
+
 
 async def test_restore_state(hass: HomeAssistant) -> None:
     """Test integration sensor state is restored correctly."""
@@ -416,13 +476,15 @@ async def test_units(hass: HomeAssistant) -> None:
     assert new_state.state == STATE_UNAVAILABLE
 
 
-async def test_device_class(hass: HomeAssistant) -> None:
+@pytest.mark.parametrize("method", ["trapezoidal", "left", "right"])
+async def test_device_class(hass: HomeAssistant, method) -> None:
     """Test integration sensor units using a power source."""
     config = {
         "sensor": {
             "platform": "integration",
             "name": "integration",
             "source": "sensor.power",
+            "method": method,
         }
     }
 
@@ -465,13 +527,15 @@ async def test_device_class(hass: HomeAssistant) -> None:
     assert state.attributes.get("device_class") == SensorDeviceClass.ENERGY
 
 
-async def test_calc_errors(hass: HomeAssistant) -> None:
+@pytest.mark.parametrize("method", ["trapezoidal", "left", "right"])
+async def test_calc_errors(hass: HomeAssistant, method) -> None:
     """Test integration sensor units using a power source."""
     config = {
         "sensor": {
             "platform": "integration",
             "name": "integration",
             "source": "sensor.power",
+            "method": method,
         }
     }
 
@@ -479,6 +543,7 @@ async def test_calc_errors(hass: HomeAssistant) -> None:
 
     entity_id = config["sensor"]["source"]
 
+    now = dt_util.utcnow()
     hass.states.async_set(entity_id, None, {})
     await hass.async_block_till_done()
 
@@ -489,19 +554,25 @@ async def test_calc_errors(hass: HomeAssistant) -> None:
     assert state.state == STATE_UNKNOWN
 
     # Moving from an unknown state to a value is a calc error and should
-    # not change the value of the Reimann sensor.
+    # not change the value of the Reimann sensor, unless the method used is "right".
-    hass.states.async_set(entity_id, 0, {"device_class": None})
+    now += timedelta(seconds=3600)
+    with patch("homeassistant.util.dt.utcnow", return_value=now):
+        hass.states.async_set(entity_id, 0, {"device_class": None})
+        await hass.async_block_till_done()
     await hass.async_block_till_done()
 
     state = hass.states.get("sensor.integration")
     assert state is not None
-    assert state.state == STATE_UNKNOWN
+    assert state.state == STATE_UNKNOWN if method != "right" else "0.000"
 
     # With the source sensor updated successfully, the Reimann sensor
     # should have a zero (known) value.
-    hass.states.async_set(entity_id, 1, {"device_class": None})
+    now += timedelta(seconds=3600)
+    with patch("homeassistant.util.dt.utcnow", return_value=now):
+        hass.states.async_set(entity_id, 1, {"device_class": None})
+        await hass.async_block_till_done()
     await hass.async_block_till_done()
 
     state = hass.states.get("sensor.integration")
     assert state is not None
-    assert round(float(state.state)) == 0
+    assert round(float(state.state)) == 0 if method != "right" else 1