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

homeassistant/components/energy/data.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 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 homeassistant.helpers import config_validation as cv, singleton, storage
 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 @@ DEVICE_CONSUMPTION_SCHEMA = vol.Schema(
 )
 
 
+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 @@ class EnergyManager:
         return {
             "energy_sources": [],
             "device_consumption": [],
+            "device_consumption_water": [],
         }
 
     async def async_update(self, update: EnergyPreferencesUpdate) -> None:
@@ -362,6 +382,7 @@ class EnergyManager:
         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 @@ class EnergyPreferencesValidation:
                 [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"] == []

tests/components/energy/test_validate.py

@@ -26,6 +26,7 @@ async def test_validation_empty_config(hass: HomeAssistant) -> None:
     assert (await validate.async_validate(hass)).as_dict() == {
         "energy_sources": [],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -80,6 +81,7 @@ async def test_validation(
     assert (await validate.async_validate(hass)).as_dict() == {
         "energy_sources": [[], []],
         "device_consumption": [[]],
+        "device_consumption_water": [],
     }
 
 
@@ -106,6 +108,7 @@ async def test_validation_device_consumption_entity_missing(
                 },
             ]
         ],
+        "device_consumption_water": [],
     }
 
 
@@ -127,6 +130,7 @@ async def test_validation_device_consumption_stat_missing(
                 }
             ]
         ],
+        "device_consumption_water": [],
     }
 
 
@@ -150,6 +154,7 @@ async def test_validation_device_consumption_entity_unavailable(
                 }
             ]
         ],
+        "device_consumption_water": [],
     }
 
 
@@ -173,6 +178,7 @@ async def test_validation_device_consumption_entity_non_numeric(
                 },
             ]
         ],
+        "device_consumption_water": [],
     }
 
 
@@ -204,6 +210,7 @@ async def test_validation_device_consumption_entity_unexpected_unit(
                 }
             ]
         ],
+        "device_consumption_water": [],
     }
 
 
@@ -227,6 +234,7 @@ async def test_validation_device_consumption_recorder_not_tracked(
                 }
             ]
         ],
+        "device_consumption_water": [],
     }
 
 
@@ -258,6 +266,7 @@ async def test_validation_device_consumption_no_last_reset(
                 }
             ]
         ],
+        "device_consumption_water": [],
     }
 
 
@@ -293,6 +302,7 @@ async def test_validation_solar(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -344,6 +354,7 @@ async def test_validation_battery(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -438,6 +449,7 @@ async def test_validation_grid(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -510,6 +522,7 @@ async def test_validation_grid_external_cost_compensation(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -585,6 +598,7 @@ async def test_validation_grid_price_not_exist(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -646,6 +660,7 @@ async def test_validation_grid_auto_cost_entity_errors(
     assert (await validate.async_validate(hass)).as_dict() == {
         "energy_sources": [[]],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -715,6 +730,7 @@ async def test_validation_grid_price_errors(
             [expected],
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -848,6 +864,7 @@ async def test_validation_gas(
             ],
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -881,6 +898,7 @@ async def test_validation_gas_no_costs_tracking(
     assert (await validate.async_validate(hass)).as_dict() == {
         "energy_sources": [[]],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -928,6 +946,7 @@ async def test_validation_grid_no_costs_tracking(
     assert (await validate.async_validate(hass)).as_dict() == {
         "energy_sources": [[]],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -1058,6 +1077,7 @@ async def test_validation_water(
             ],
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -1091,4 +1111,138 @@ async def test_validation_water_no_costs_tracking(
     assert (await validate.async_validate(hass)).as_dict() == {
         "energy_sources": [[]],
         "device_consumption": [],
+        "device_consumption_water": [],
+    }
+
+
+async def test_validation_device_consumption_water_entity_missing(
+    hass: HomeAssistant, mock_energy_manager
+) -> None:
+    """Test validating missing entity for water device."""
+    await mock_energy_manager.async_update(
+        {"device_consumption_water": [{"stat_consumption": "sensor.not_exist"}]}
+    )
+    assert (await validate.async_validate(hass)).as_dict() == {
+        "energy_sources": [],
+        "device_consumption": [],
+        "device_consumption_water": [
+            [
+                {
+                    "type": "statistics_not_defined",
+                    "affected_entities": {("sensor.not_exist", None)},
+                    "translation_placeholders": None,
+                },
+                {
+                    "type": "entity_not_defined",
+                    "affected_entities": {("sensor.not_exist", None)},
+                    "translation_placeholders": None,
+                },
+            ]
+        ],
+    }
+
+
+async def test_validation_device_consumption_water_entity_unexpected_unit(
+    hass: HomeAssistant, mock_energy_manager, mock_get_metadata
+) -> None:
+    """Test validating water device with unexpected unit."""
+    await mock_energy_manager.async_update(
+        {"device_consumption_water": [{"stat_consumption": "sensor.unexpected_unit"}]}
+    )
+    hass.states.async_set(
+        "sensor.unexpected_unit",
+        "10.10",
+        {
+            "device_class": "water",
+            "unit_of_measurement": "beers",
+            "state_class": "total_increasing",
+        },
+    )
+
+    assert (await validate.async_validate(hass)).as_dict() == {
+        "energy_sources": [],
+        "device_consumption": [],
+        "device_consumption_water": [
+            [
+                {
+                    "type": "entity_unexpected_unit_water",
+                    "affected_entities": {("sensor.unexpected_unit", "beers")},
+                    "translation_placeholders": {
+                        "water_units": "CCF, ft³, m³, gal, L, MCF"
+                    },
+                }
+            ]
+        ],
+    }
+
+
+async def test_validation_device_consumption_water_valid_units(
+    hass: HomeAssistant, mock_energy_manager, mock_get_metadata
+) -> None:
+    """Test validating water device with valid water units."""
+    await mock_energy_manager.async_update(
+        {
+            "device_consumption_water": [
+                {"stat_consumption": "sensor.water_m3"},
+                {"stat_consumption": "sensor.water_l"},
+                {"stat_consumption": "sensor.water_gal"},
+            ]
+        }
+    )
+    hass.states.async_set(
+        "sensor.water_m3",
+        "10.10",
+        {
+            "device_class": "water",
+            "unit_of_measurement": "m³",
+            "state_class": "total_increasing",
+        },
+    )
+    hass.states.async_set(
+        "sensor.water_l",
+        "1000.0",
+        {
+            "device_class": "water",
+            "unit_of_measurement": "L",
+            "state_class": "total_increasing",
+        },
+    )
+    hass.states.async_set(
+        "sensor.water_gal",
+        "100.0",
+        {
+            "device_class": "water",
+            "unit_of_measurement": "gal",
+            "state_class": "total_increasing",
+        },
+    )
+
+    assert (await validate.async_validate(hass)).as_dict() == {
+        "energy_sources": [],
+        "device_consumption": [],
+        "device_consumption_water": [[], [], []],
+    }
+
+
+async def test_validation_device_consumption_water_recorder_not_tracked(
+    hass: HomeAssistant, mock_energy_manager, mock_is_entity_recorded, mock_get_metadata
+) -> None:
+    """Test validating water device based on untracked entity."""
+    mock_is_entity_recorded["sensor.not_recorded"] = False
+    await mock_energy_manager.async_update(
+        {"device_consumption_water": [{"stat_consumption": "sensor.not_recorded"}]}
+    )
+
+    assert (await validate.async_validate(hass)).as_dict() == {
+        "energy_sources": [],
+        "device_consumption": [],
+        "device_consumption_water": [
+            [
+                {
+                    "type": "recorder_untracked",
+                    "affected_entities": {("sensor.not_recorded", None)},
+                    "translation_placeholders": None,
+                }
+            ]
+        ],
     }

tests/components/energy/test_validate_power.py

@@ -53,6 +53,7 @@ async def test_validation_grid_power_valid(
     assert result.as_dict() == {
         "energy_sources": [[]],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -99,6 +100,7 @@ async def test_validation_grid_power_wrong_unit(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -145,6 +147,7 @@ async def test_validation_grid_power_wrong_state_class(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -187,6 +190,7 @@ async def test_validation_grid_power_entity_missing(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -225,6 +229,7 @@ async def test_validation_grid_power_entity_unavailable(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -271,6 +276,7 @@ async def test_validation_grid_power_entity_non_numeric(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -319,6 +325,7 @@ async def test_validation_grid_power_wrong_device_class(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -369,6 +376,7 @@ async def test_validation_grid_power_different_units(
     assert result.as_dict() == {
         "energy_sources": [[]],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -408,6 +416,7 @@ async def test_validation_grid_power_external_statistics(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }
 
 
@@ -447,4 +456,5 @@ async def test_validation_grid_power_recorder_untracked(
             ]
         ],
         "device_consumption": [],
+        "device_consumption_water": [],
     }

tests/components/energy/test_websocket_api.py

@@ -165,6 +165,12 @@ async def test_save_preferences(
                 "stat_rate": "sensor.some_device_power",
             }
         ],
+        "device_consumption_water": [
+            {
+                "stat_consumption": "sensor.water_meter",
+                "name": "Water Meter",
+            }
+        ],
     }
 
     await client.send_json({"id": 6, "type": "energy/save_prefs", **new_prefs})
@@ -290,6 +296,7 @@ async def test_validate(
     assert msg["result"] == {
         "energy_sources": [],
         "device_consumption": [],
+        "device_consumption_water": [],
     }