Add statistics importing for Anglian Water (#157757)

homeassistant/components/anglian_water/coordinator.py

@@ -4,13 +4,28 @@ from __future__ import annotations
 
 from datetime import timedelta
 import logging
+from typing import Any
 
 from pyanglianwater import AnglianWater
 from pyanglianwater.exceptions import ExpiredAccessTokenError, UnknownEndpointError
 
+from homeassistant.components.recorder import get_instance
+from homeassistant.components.recorder.models import (
+    StatisticData,
+    StatisticMeanType,
+    StatisticMetaData,
+)
+from homeassistant.components.recorder.statistics import (
+    async_add_external_statistics,
+    get_last_statistics,
+    statistics_during_period,
+)
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import UnitOfVolume
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+from homeassistant.util import dt as dt_util
+from homeassistant.util.unit_conversion import VolumeConverter
 
 from .const import CONF_ACCOUNT_NUMBER, DOMAIN
 
@@ -44,6 +59,107 @@ class AnglianWaterUpdateCoordinator(DataUpdateCoordinator[None]):
     async def _async_update_data(self) -> None:
         """Update data from Anglian Water's API."""
         try:
-            return await self.api.update(self.config_entry.data[CONF_ACCOUNT_NUMBER])
+            await self.api.update(self.config_entry.data[CONF_ACCOUNT_NUMBER])
+            await self._insert_statistics()
         except (ExpiredAccessTokenError, UnknownEndpointError) as err:
             raise UpdateFailed from err
+
+    async def _insert_statistics(self) -> None:
+        """Insert statistics for water meters into Home Assistant."""
+        for meter in self.api.meters.values():
+            id_prefix = (
+                f"{self.config_entry.data[CONF_ACCOUNT_NUMBER]}_{meter.serial_number}"
+            )
+            usage_statistic_id = f"{DOMAIN}:{id_prefix}_usage".lower()
+            _LOGGER.debug("Updating statistics for meter %s", meter.serial_number)
+            name_prefix = (
+                f"Anglian Water {self.config_entry.data[CONF_ACCOUNT_NUMBER]} "
+                f"{meter.serial_number}"
+            )
+            usage_metadata = StatisticMetaData(
+                mean_type=StatisticMeanType.NONE,
+                has_sum=True,
+                name=f"{name_prefix} Usage",
+                source=DOMAIN,
+                statistic_id=usage_statistic_id,
+                unit_class=VolumeConverter.UNIT_CLASS,
+                unit_of_measurement=UnitOfVolume.CUBIC_METERS,
+            )
+            last_stat = await get_instance(self.hass).async_add_executor_job(
+                get_last_statistics, self.hass, 1, usage_statistic_id, True, set()
+            )
+            if not last_stat:
+                _LOGGER.debug("Updating statistics for the first time")
+                usage_sum = 0.0
+                last_stats_time = None
+            else:
+                if not meter.readings or len(meter.readings) == 0:
+                    _LOGGER.debug("No recent usage statistics found, skipping update")
+                    continue
+                # Anglian Water stats are hourly, the read_at time is the time that the meter took the reading
+                # We remove 1 hour from this so that the data is shown in the correct hour on the dashboards
+                parsed_read_at = dt_util.parse_datetime(meter.readings[0]["read_at"])
+                if not parsed_read_at:
+                    _LOGGER.debug(
+                        "Could not parse read_at time %s, skipping update",
+                        meter.readings[0]["read_at"],
+                    )
+                    continue
+                start = dt_util.as_local(parsed_read_at) - timedelta(hours=1)
+                _LOGGER.debug("Getting statistics at %s", start)
+                for end in (start + timedelta(seconds=1), None):
+                    stats = await get_instance(self.hass).async_add_executor_job(
+                        statistics_during_period,
+                        self.hass,
+                        start,
+                        end,
+                        {
+                            usage_statistic_id,
+                        },
+                        "hour",
+                        None,
+                        {"sum"},
+                    )
+                    if stats:
+                        break
+                    if end:
+                        _LOGGER.debug(
+                            "Not found, trying to find oldest statistic after %s",
+                            start,
+                        )
+                assert stats
+
+                def _safe_get_sum(records: list[Any]) -> float:
+                    if records and "sum" in records[0]:
+                        return float(records[0]["sum"])
+                    return 0.0
+
+                usage_sum = _safe_get_sum(stats.get(usage_statistic_id, []))
+                last_stats_time = stats[usage_statistic_id][0]["start"]
+
+            usage_statistics = []
+
+            for read in meter.readings:
+                parsed_read_at = dt_util.parse_datetime(read["read_at"])
+                if not parsed_read_at:
+                    _LOGGER.debug(
+                        "Could not parse read_at time %s, skipping reading",
+                        read["read_at"],
+                    )
+                    continue
+                start = dt_util.as_local(parsed_read_at) - timedelta(hours=1)
+                if last_stats_time is not None and start.timestamp() <= last_stats_time:
+                    continue
+                usage_state = max(0, read["consumption"] / 1000)
+                usage_sum = max(0, read["read"])
+                usage_statistics.append(
+                    StatisticData(
+                        start=start,
+                        state=usage_state,
+                        sum=usage_sum,
+                    )
+                )
+            _LOGGER.debug(
+                "Adding %s statistics for %s", len(usage_statistics), usage_statistic_id
+            )
+            async_add_external_statistics(self.hass, usage_metadata, usage_statistics)

homeassistant/components/anglian_water/manifest.json

@@ -1,6 +1,7 @@
 {
   "domain": "anglian_water",
   "name": "Anglian Water",
+  "after_dependencies": ["recorder"],
   "codeowners": ["@pantherale0"],
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/anglian_water",

tests/components/anglian_water/conftest.py

@@ -38,6 +38,11 @@ def mock_smart_meter() -> SmartMeter:
     mock.latest_read = 50
     mock.yesterday_water_cost = 0.5
     mock.yesterday_sewerage_cost = 0.5
+    mock.readings = [
+        {"read_at": "2024-06-01T12:00:00Z", "consumption": 10, "read": 10},
+        {"read_at": "2024-06-01T13:00:00Z", "consumption": 15, "read": 25},
+        {"read_at": "2024-06-01T14:00:00Z", "consumption": 25, "read": 50},
+    ]
     return mock
 
 

tests/components/anglian_water/snapshots/test_coordinator.ambr

@@ -0,0 +1,55 @@
+# serializer version: 1
+# name: test_coordinator_first_run
+  defaultdict({
+    'anglian_water:12345678_testsn_usage': list([
+      dict({
+        'end': 1717243200.0,
+        'start': 1717239600.0,
+        'state': 0.01,
+        'sum': 10.0,
+      }),
+      dict({
+        'end': 1717246800.0,
+        'start': 1717243200.0,
+        'state': 0.015,
+        'sum': 25.0,
+      }),
+      dict({
+        'end': 1717250400.0,
+        'start': 1717246800.0,
+        'state': 0.025,
+        'sum': 50.0,
+      }),
+    ]),
+  })
+# ---
+# name: test_coordinator_subsequent_run
+  defaultdict({
+    'anglian_water:12345678_testsn_usage': list([
+      dict({
+        'end': 1717243200.0,
+        'start': 1717239600.0,
+        'state': 0.01,
+        'sum': 10.0,
+      }),
+      dict({
+        'end': 1717246800.0,
+        'start': 1717243200.0,
+        'state': 0.015,
+        'sum': 25.0,
+      }),
+      dict({
+        'end': 1717250400.0,
+        'start': 1717246800.0,
+        'state': 0.035,
+        'sum': 70.0,
+      }),
+      dict({
+        'end': 1717254000.0,
+        'start': 1717250400.0,
+        'state': 0.02,
+        'sum': 90.0,
+      }),
+    ]),
+  })
+# ---

tests/components/anglian_water/test_coordinator.py

@@ -0,0 +1,164 @@
+"""Tests for the Anglian Water coordinator."""
+
+from unittest.mock import AsyncMock
+
+from pyanglianwater.meter import SmartMeter
+import pytest
+from syrupy.assertion import SnapshotAssertion
+
+from homeassistant.components.anglian_water.coordinator import (
+    AnglianWaterUpdateCoordinator,
+)
+from homeassistant.components.recorder import Recorder
+from homeassistant.components.recorder.statistics import (
+    get_last_statistics,
+    statistics_during_period,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.util import dt as dt_util
+
+from .const import ACCOUNT_NUMBER
+
+from tests.common import MockConfigEntry
+from tests.components.recorder.common import async_wait_recording_done
+
+
+async def test_coordinator_first_run(
+    recorder_mock: Recorder,
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_anglian_water_client: AsyncMock,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test the coordinator on its first run with no existing statistics."""
+    coordinator = AnglianWaterUpdateCoordinator(
+        hass, mock_anglian_water_client, mock_config_entry
+    )
+    await coordinator._async_update_data()
+    await async_wait_recording_done(hass)
+
+    stats = await hass.async_add_executor_job(
+        statistics_during_period,
+        hass,
+        dt_util.utc_from_timestamp(0),
+        None,
+        {
+            f"anglian_water:{ACCOUNT_NUMBER}_testsn_usage",
+        },
+        "hour",
+        None,
+        {"state", "sum"},
+    )
+    assert stats == snapshot
+
+
+async def test_coordinator_subsequent_run(
+    recorder_mock: Recorder,
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_smart_meter: SmartMeter,
+    mock_anglian_water_client: AsyncMock,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test the coordinator correctly updates statistics on subsequent runs."""
+    # 1st run
+    coordinator = AnglianWaterUpdateCoordinator(
+        hass, mock_anglian_water_client, mock_config_entry
+    )
+    await coordinator._async_update_data()
+    await async_wait_recording_done(hass)
+
+    # 2nd run with an updated reading for one read and a new read added.
+    mock_smart_meter.readings[-1] = {
+        "read_at": "2024-06-01T14:00:00Z",
+        "consumption": 35,
+        "read": 70,
+    }
+    mock_smart_meter.readings.append(
+        {"read_at": "2024-06-01T15:00:00Z", "consumption": 20, "read": 90}
+    )
+    await coordinator._async_update_data()
+    await async_wait_recording_done(hass)
+
+    # Check all stats
+    stats = await hass.async_add_executor_job(
+        statistics_during_period,
+        hass,
+        dt_util.utc_from_timestamp(0),
+        None,
+        {
+            f"anglian_water:{ACCOUNT_NUMBER}_testsn_usage",
+        },
+        "hour",
+        None,
+        {"state", "sum"},
+    )
+    assert stats == snapshot
+
+
+async def test_coordinator_subsequent_run_no_energy_data(
+    recorder_mock: Recorder,
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_smart_meter: SmartMeter,
+    mock_anglian_water_client: AsyncMock,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test the coordinator handles no recent usage/cost data."""
+    # 1st run
+    coordinator = AnglianWaterUpdateCoordinator(
+        hass, mock_anglian_water_client, mock_config_entry
+    )
+    await coordinator._async_update_data()
+    await async_wait_recording_done(hass)
+
+    # 2nd run with no readings
+    mock_smart_meter.readings = []
+    await coordinator._async_update_data()
+
+    assert "No recent usage statistics found, skipping update" in caplog.text
+    # Verify no new stats were added by checking the sum remains 50
+    statistic_id = f"anglian_water:{ACCOUNT_NUMBER}_testsn_usage"
+    stats = await hass.async_add_executor_job(
+        get_last_statistics, hass, 1, statistic_id, True, {"sum"}
+    )
+    assert stats[statistic_id][0]["sum"] == 50
+
+
+async def test_coordinator_invalid_readings(
+    recorder_mock: Recorder,
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_smart_meter: SmartMeter,
+    mock_anglian_water_client: AsyncMock,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test the coordinator handles bad data / invalid readings correctly."""
+    coordinator = AnglianWaterUpdateCoordinator(
+        hass, mock_anglian_water_client, mock_config_entry
+    )
+    await coordinator._async_update_data()
+    await async_wait_recording_done(hass)
+
+    # Test that an invalid read_at on the first reading skips the entire update
+    mock_smart_meter.readings = [
+        {"read_at": "invalid-date-format", "consumption": 10, "read": 10},
+    ]
+    await coordinator._async_update_data()
+
+    assert (
+        "Could not parse read_at time invalid-date-format, skipping update"
+        in caplog.text
+    )
+
+    # Test that individual invalid readings are skipped
+    mock_smart_meter.readings = [
+        {"read_at": "2024-06-01T12:00:00Z", "consumption": 10, "read": 10},
+        {"read_at": "also-invalid-date", "consumption": 15, "read": 25},
+    ]
+    await coordinator._async_update_data()
+
+    assert (
+        "Could not parse read_at time also-invalid-date, skipping reading"
+        in caplog.text
+    )

tests/components/anglian_water/test_sensor.py

@@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, patch
 
 from syrupy.assertion import SnapshotAssertion
 
+from homeassistant.components.recorder import Recorder
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
@@ -14,6 +15,7 @@ from tests.common import MockConfigEntry, snapshot_platform
 
 
 async def test_sensor(
+    recorder_mock: Recorder,
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,
     mock_anglian_water_client: AsyncMock,