Suez water: fetch historical data in statistics

jb101010-2 · jb101010-2 · commit 58d5f8db5038 · 2024-11-21T15:10:26.000+01:00
diff --git a/homeassistant/components/suez_water/coordinator.py b/homeassistant/components/suez_water/coordinator.py
@@ -2,13 +2,26 @@
 
 from collections.abc import Mapping
 from dataclasses import dataclass
-from datetime import date
+from datetime import date, datetime, time
 from typing import Any
+from zoneinfo import ZoneInfo
 
-from pysuez import PySuezError, SuezClient
+from pysuez import DayDataResult, PySuezError, SuezClient
 
+from homeassistant.components.recorder import get_instance
+from homeassistant.components.recorder.models import StatisticData, StatisticMetaData
+from homeassistant.components.recorder.statistics import (
+    StatisticsRow,
+    async_add_external_statistics,
+    get_last_statistics,
+)
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
+from homeassistant.const import (
+    CONF_PASSWORD,
+    CONF_USERNAME,
+    CURRENCY_EURO,
+    UnitOfVolume,
+)
 from homeassistant.core import _LOGGER, HomeAssistant
 from homeassistant.exceptions import ConfigEntryError
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
@@ -53,6 +66,12 @@ def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None:
             always_update=True,
             config_entry=config_entry,
         )
+        self._counter_id = self.config_entry.data[CONF_COUNTER_ID]
+        self._cost_statistic_id = f"{DOMAIN}:{self._counter_id}_water_cost_statistics"
+        self._water_statistic_id = (
+            f"{DOMAIN}:{self._counter_id}_water_consumption_statistics"
+        )
+        self.config_entry.async_on_unload(self._clear_statistics)
 
     async def _async_setup(self) -> None:
         self._suez_client = SuezClient(
@@ -79,10 +98,142 @@ async def _async_update_data(self) -> SuezWaterData:
                 },
                 price=(await self._suez_client.get_price()).price,
             )
+            await self._update_statistics(data.price)
         except PySuezError as err:
             _LOGGER.exception(err)
             raise UpdateFailed(
                 f"Suez coordinator error communicating with API: {err}"
             ) from err
         _LOGGER.debug("Successfully fetched suez data")
         return data
+
+    async def _update_statistics(self, current_price: float) -> None:
+        """Update daily statistics."""
+        _LOGGER.debug("Updating statistics for %s", self._water_statistic_id)
+
+        water_last_stat = await self._get_last_stat(self._water_statistic_id)
+        cost_last_stat = await self._get_last_stat(self._cost_statistic_id)
+
+        consumption_sum = 0.0
+        cost_sum = 0.0
+        last_stats = None
+
+        if water_last_stat is not None:
+            last_stats = datetime.fromtimestamp(water_last_stat["start"]).date()
+            if water_last_stat["sum"] is not None:
+                consumption_sum = water_last_stat["sum"]
+        if cost_last_stat is not None:
+            if cost_last_stat["sum"] is not None and cost_last_stat["sum"] is not None:
+                cost_sum = cost_last_stat["sum"]
+
+        _LOGGER.debug(
+            "Updating suez stat since %s for %s",
+            str(last_stats),
+            water_last_stat,
+        )
+        usage = await self._suez_client.fetch_all_daily_data(
+            since=last_stats,
+        )
+        if usage is None or len(usage) <= 0:
+            _LOGGER.debug("No recent usage data. Skipping update")
+            return
+        _LOGGER.debug("fetched data: %s", len(usage))
+
+        consumption_statistics, cost_statistics = self._build_statistics(
+            current_price, consumption_sum, cost_sum, last_stats, usage
+        )
+
+        self._persist_statistics(consumption_statistics, cost_statistics)
+
+    def _build_statistics(
+        self,
+        current_price: float,
+        consumption_sum: float,
+        cost_sum: float,
+        last_stats: date | None,
+        usage: list[DayDataResult],
+    ) -> tuple[list[StatisticData], list[StatisticData]]:
+        """Build statistics data from fetched data."""
+        consumption_statistics = []
+        cost_statistics = []
+
+        for data in usage:
+            if last_stats is not None and data.date <= last_stats:
+                continue
+            consumption_date = datetime.combine(
+                data.date, time(0, 0, 0, 0), ZoneInfo("Europe/Paris")
+            )
+
+            consumption_sum += data.day_consumption
+            consumption_statistics.append(
+                StatisticData(
+                    start=consumption_date,
+                    state=data.day_consumption,
+                    sum=consumption_sum,
+                )
+            )
+            day_cost = (data.day_consumption / 1000) * current_price
+            cost_sum += day_cost
+            cost_statistics.append(
+                StatisticData(
+                    start=consumption_date,
+                    state=day_cost,
+                    sum=cost_sum,
+                )
+            )
+
+        return consumption_statistics, cost_statistics
+
+    def _persist_statistics(
+        self,
+        consumption_statistics: list[StatisticData],
+        cost_statistics: list[StatisticData],
+    ) -> None:
+        """Persist given statistics in recorder."""
+        consumption_metadata = self._get_statistics_metadata(
+            id=self._water_statistic_id, name="Consumption", unit=UnitOfVolume.LITERS
+        )
+        cost_metadata = self._get_statistics_metadata(
+            id=self._cost_statistic_id, name="Cost", unit=CURRENCY_EURO
+        )
+
+        _LOGGER.debug(
+            "Adding %s statistics for %s",
+            len(consumption_statistics),
+            self._water_statistic_id,
+        )
+        async_add_external_statistics(
+            self.hass, consumption_metadata, consumption_statistics
+        )
+        async_add_external_statistics(self.hass, cost_metadata, cost_statistics)
+
+        _LOGGER.debug("Updated statistics for %s", self._water_statistic_id)
+
+    def _get_statistics_metadata(
+        self, id: str, name: str, unit: str
+    ) -> StatisticMetaData:
+        """Build statistics metadata for requested configuration."""
+        return StatisticMetaData(
+            has_mean=False,
+            has_sum=True,
+            name=f"Suez water {name} {self._counter_id}",
+            source=DOMAIN,
+            statistic_id=id,
+            unit_of_measurement=unit,
+        )
+
+    async def _get_last_stat(self, id: str) -> StatisticsRow | None:
+        """Find last registered statistics of given id."""
+        last_stat = await get_instance(self.hass).async_add_executor_job(
+            get_last_statistics, self.hass, 1, id, True, {"sum"}
+        )
+        if last_stat is None or len(last_stat) == 0:
+            return None
+        return last_stat[id][0]
+
+    def _clear_statistics(self) -> None:
+        """Clear suez water statistics."""
+        instance = get_instance(self.hass)
+        instance.async_clear_statistics(
+            [self._water_statistic_id, self._cost_statistic_id]
+        )
diff --git a/homeassistant/components/suez_water/manifest.json b/homeassistant/components/suez_water/manifest.json
@@ -3,6 +3,7 @@
   "name": "Suez Water",
   "codeowners": ["@ooii", "@jb101010-2"],
   "config_flow": true,
+  "dependencies": ["recorder"],
   "documentation": "https://www.home-assistant.io/integrations/suez_water",
   "iot_class": "cloud_polling",
   "loggers": ["pysuez", "regex"],
diff --git a/tests/components/suez_water/conftest.py b/tests/components/suez_water/conftest.py
@@ -7,17 +7,26 @@
 from pysuez.const import ATTRIBUTION
 import pytest
 
+from homeassistant.components.recorder import Recorder
 from homeassistant.components.suez_water.const import DOMAIN
 
 from tests.common import MockConfigEntry
+from tests.conftest import RecorderInstanceGenerator
 
 MOCK_DATA = {
     "username": "test-username",
     "password": "test-password",
-    "counter_id": "test-counter",
+    "counter_id": "123456",
 }
 
 
+@pytest.fixture
+async def mock_recorder_before_hass(
+    async_test_recorder: RecorderInstanceGenerator,
+) -> None:
+    """Set up recorder."""
+
+
 @pytest.fixture
 def mock_config_entry() -> MockConfigEntry:
     """Create mock config_entry needed by suez_water integration."""
@@ -30,7 +39,7 @@ def mock_config_entry() -> MockConfigEntry:
 
 
 @pytest.fixture
-def mock_setup_entry() -> Generator[AsyncMock]:
+def mock_setup_entry(recorder_mock: Recorder) -> Generator[AsyncMock]:
     """Override async_setup_entry."""
     with patch(
         "homeassistant.components.suez_water.async_setup_entry", return_value=True
@@ -39,7 +48,7 @@ def mock_setup_entry() -> Generator[AsyncMock]:
 
 
 @pytest.fixture(name="suez_client")
-def mock_suez_client() -> Generator[AsyncMock]:
+def mock_suez_client(recorder_mock: Recorder) -> Generator[AsyncMock]:
     """Create mock for suez_water external api."""
     with (
         patch(
diff --git a/tests/components/suez_water/snapshots/test_sensor.ambr b/tests/components/suez_water/snapshots/test_sensor.ambr
@@ -28,7 +28,7 @@
     'previous_unique_id': None,
     'supported_features': 0,
     'translation_key': 'water_price',
-    'unique_id': 'test-counter_water_price',
+    'unique_id': '123456_water_price',
     'unit_of_measurement': '€',
   })
 # ---
@@ -77,7 +77,7 @@
     'previous_unique_id': None,
     'supported_features': 0,
     'translation_key': 'water_usage_yesterday',
-    'unique_id': 'test-counter_water_usage_yesterday',
+    'unique_id': '123456_water_usage_yesterday',
     'unit_of_measurement': <UnitOfVolume.LITERS: 'L'>,
   })
 # ---
diff --git a/tests/components/suez_water/test_config_flow.py b/tests/components/suez_water/test_config_flow.py
@@ -6,6 +6,7 @@
 import pytest
 
 from homeassistant import config_entries
+from homeassistant.components.recorder import Recorder
 from homeassistant.components.suez_water.const import CONF_COUNTER_ID, DOMAIN
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
@@ -69,7 +70,9 @@ async def test_form_invalid_auth(
     assert len(mock_setup_entry.mock_calls) == 1
 
 
-async def test_form_already_configured(hass: HomeAssistant) -> None:
+async def test_form_already_configured(
+    hass: HomeAssistant, recorder_mock: Recorder
+) -> None:
     """Test we abort when entry is already configured."""
 
     entry = MockConfigEntry(
diff --git a/tests/components/suez_water/test_sensor.py b/tests/components/suez_water/test_sensor.py
@@ -26,7 +26,13 @@ async def test_sensors_valid_state(
     entity_registry: er.EntityRegistry,
 ) -> None:
     """Test that suez_water sensor is loaded and in a valid state."""
-    with patch("homeassistant.components.suez_water.PLATFORMS", [Platform.SENSOR]):
+    with (
+        patch("homeassistant.components.suez_water.PLATFORMS", [Platform.SENSOR]),
+        patch(
+            "homeassistant.components.suez_water.coordinator.SuezWaterCoordinator._update_statistics",
+            return_value=None,
+        ),
+    ):
         await setup_integration(hass, mock_config_entry)
 
     assert mock_config_entry.state is ConfigEntryState.LOADED
@@ -43,25 +49,31 @@ async def test_sensors_failed_update(
 ) -> None:
     """Test that suez_water sensor reflect failure when api fails."""
 
-    await setup_integration(hass, mock_config_entry)
+    with (
+        patch(
+            "homeassistant.components.suez_water.coordinator.SuezWaterCoordinator._update_statistics",
+            return_value=None,
+        ),
+    ):
+        await setup_integration(hass, mock_config_entry)
 
-    assert mock_config_entry.state is ConfigEntryState.LOADED
+        assert mock_config_entry.state is ConfigEntryState.LOADED
 
-    entity_ids = await hass.async_add_executor_job(hass.states.entity_ids)
-    assert len(entity_ids) == 2
+        entity_ids = await hass.async_add_executor_job(hass.states.entity_ids)
+        assert len(entity_ids) == 2
 
-    for entity in entity_ids:
-        state = hass.states.get(entity)
-        assert entity
-        assert state.state != STATE_UNAVAILABLE
+        for entity in entity_ids:
+            state = hass.states.get(entity)
+            assert entity
+            assert state.state != STATE_UNAVAILABLE
 
-    getattr(suez_client, method).side_effect = PySuezError("Should fail to update")
+        getattr(suez_client, method).side_effect = PySuezError("Should fail to update")
 
-    freezer.tick(DATA_REFRESH_INTERVAL)
-    async_fire_time_changed(hass)
-    await hass.async_block_till_done(True)
+        freezer.tick(DATA_REFRESH_INTERVAL)
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done(True)
 
-    for entity in entity_ids:
-        state = hass.states.get(entity)
-        assert entity
-        assert state.state == STATE_UNAVAILABLE
+        for entity in entity_ids:
+            state = hass.states.get(entity)
+            assert entity
+            assert state.state == STATE_UNAVAILABLE
diff --git a/tests/components/suez_water/test_statistics.py b/tests/components/suez_water/test_statistics.py
