Add support for downstream water meters in energy dashboard (home-assistant#155927)

Author: MindFreeze

homeassistant/components/energy/data.py

@@ -5,7 +5,7 @@
 import asyncio
 from collections import Counter
 from collections.abc import Awaitable, Callable
-from typing import Literal, NotRequired, TypedDict
+from typing import Any, Literal, NotRequired, TypedDict
 
 import voluptuous as vol
 
@@ -15,6 +15,7 @@
 from .const import DOMAIN
 
 STORAGE_VERSION = 1
+STORAGE_MINOR_VERSION = 2
 STORAGE_KEY = DOMAIN
 
 
@@ -164,6 +165,7 @@ class EnergyPreferences(TypedDict):
 
     energy_sources: list[SourceType]
     device_consumption: list[DeviceConsumption]
+    device_consumption_water: NotRequired[list[DeviceConsumption]]
 
 
 class EnergyPreferencesUpdate(EnergyPreferences, total=False):
@@ -328,14 +330,31 @@ def check_type_limits(value: list[SourceType]) -> list[SourceType]:
 )
 
 
+class _EnergyPreferencesStore(storage.Store[EnergyPreferences]):
+    """Energy preferences store with migration support."""
+
+    async def _async_migrate_func(
+        self,
+        old_major_version: int,
+        old_minor_version: int,
+        old_data: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Migrate to the new version."""
+        data = old_data
+        if old_major_version == 1 and old_minor_version < 2:
+            # Add device_consumption_water field if it doesn't exist
+            data.setdefault("device_consumption_water", [])
+        return data
+
+
 class EnergyManager:
     """Manage the instance energy prefs."""
 
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize energy manager."""
         self._hass = hass
-        self._store = storage.Store[EnergyPreferences](
-            hass, STORAGE_VERSION, STORAGE_KEY
+        self._store = _EnergyPreferencesStore(
+            hass, STORAGE_VERSION, STORAGE_KEY, minor_version=STORAGE_MINOR_VERSION
         )
         self.data: EnergyPreferences | None = None
         self._update_listeners: list[Callable[[], Awaitable]] = []
@@ -350,6 +369,7 @@ def default_preferences() -> EnergyPreferences:
         return {
             "energy_sources": [],
             "device_consumption": [],
+            "device_consumption_water": [],
         }
 
     async def async_update(self, update: EnergyPreferencesUpdate) -> None:
@@ -362,6 +382,7 @@ async def async_update(self, update: EnergyPreferencesUpdate) -> None:
         for key in (
             "energy_sources",
             "device_consumption",
+            "device_consumption_water",
         ):
             if key in update:
                 data[key] = update[key]

homeassistant/components/energy/validate.py

@@ -153,6 +153,9 @@ class EnergyPreferencesValidation:
 
     energy_sources: list[ValidationIssues] = dataclasses.field(default_factory=list)
     device_consumption: list[ValidationIssues] = dataclasses.field(default_factory=list)
+    device_consumption_water: list[ValidationIssues] = dataclasses.field(
+        default_factory=list
+    )
 
     def as_dict(self) -> dict:
         """Return dictionary version."""
@@ -165,6 +168,10 @@ def as_dict(self) -> dict:
                 [dataclasses.asdict(issue) for issue in issues.issues.values()]
                 for issues in self.device_consumption
             ],
+            "device_consumption_water": [
+                [dataclasses.asdict(issue) for issue in issues.issues.values()]
+                for issues in self.device_consumption_water
+            ],
         }
 
 
@@ -742,6 +749,23 @@ async def async_validate(hass: HomeAssistant) -> EnergyPreferencesValidation:
             )
         )
 
+    for device in manager.data.get("device_consumption_water", []):
+        device_result = ValidationIssues()
+        result.device_consumption_water.append(device_result)
+        wanted_statistics_metadata.add(device["stat_consumption"])
+        validate_calls.append(
+            functools.partial(
+                _async_validate_usage_stat,
+                hass,
+                statistics_metadata,
+                device["stat_consumption"],
+                WATER_USAGE_DEVICE_CLASSES,
+                WATER_USAGE_UNITS,
+                WATER_UNIT_ERROR,
+                device_result,
+            )
+        )
+
     # Fetch the needed statistics metadata
     statistics_metadata.update(
         await recorder.get_instance(hass).async_add_executor_job(

homeassistant/components/energy/websocket_api.py

@@ -129,6 +129,7 @@ def ws_get_prefs(
         vol.Required("type"): "energy/save_prefs",
         vol.Optional("energy_sources"): ENERGY_SOURCE_SCHEMA,
         vol.Optional("device_consumption"): [DEVICE_CONSUMPTION_SCHEMA],
+        vol.Optional("device_consumption_water"): [DEVICE_CONSUMPTION_SCHEMA],
     }
 )
 @websocket_api.async_response

tests/components/energy/test_data.py

@@ -0,0 +1,74 @@
+"""Test energy data storage and migration."""
+
+from homeassistant.components.energy.data import EnergyManager
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import storage
+
+
+async def test_energy_preferences_no_migration_needed(hass: HomeAssistant) -> None:
+    """Test that new data format doesn't get migrated."""
+    # Create new format data (already has device_consumption_water field)
+    new_data = {
+        "energy_sources": [],
+        "device_consumption": [],
+        "device_consumption_water": [
+            {"stat_consumption": "sensor.water_meter", "name": "Water heater"}
+        ],
+    }
+
+    # Save data that already has the new field
+    old_store = storage.Store(hass, 1, "energy", minor_version=1)
+    await old_store.async_save(new_data)
+
+    # Load it with manager
+    manager = EnergyManager(hass)
+    await manager.async_initialize()
+
+    # Verify the data is unchanged
+    assert manager.data is not None
+    assert manager.data["device_consumption_water"] == [
+        {"stat_consumption": "sensor.water_meter", "name": "Water heater"}
+    ]
+
+
+async def test_energy_preferences_default(hass: HomeAssistant) -> None:
+    """Test default preferences include device_consumption_water."""
+    defaults = EnergyManager.default_preferences()
+
+    assert "energy_sources" in defaults
+    assert "device_consumption" in defaults
+    assert "device_consumption_water" in defaults
+    assert defaults["device_consumption_water"] == []
+
+
+async def test_energy_preferences_empty_store(hass: HomeAssistant) -> None:
+    """Test loading with no existing data."""
+    manager = EnergyManager(hass)
+    await manager.async_initialize()
+
+    # Verify data is None when no existing data
+    assert manager.data is None
+
+
+async def test_energy_preferences_migration_from_old_version(
+    hass: HomeAssistant,
+) -> None:
+    """Test that device_consumption_water is added when migrating from v1.1 to v1.2."""
+    # Create version 1.1 data without device_consumption_water (old version)
+    old_data = {
+        "energy_sources": [],
+        "device_consumption": [],
+    }
+
+    # Save with old version (1.1) - migration will run to upgrade to 1.2
+    old_store = storage.Store(hass, 1, "energy", minor_version=1)
+    await old_store.async_save(old_data)
+
+    # Load with manager - should trigger migration
+    manager = EnergyManager(hass)
+    await manager.async_initialize()
+
+    # Verify the field was added by migration
+    assert manager.data is not None
+    assert "device_consumption_water" in manager.data
+    assert manager.data["device_consumption_water"] == []