Add support to the energy integration for tracking water usage (#80888)

homeassistant/components/energy/data.py

@@ -87,13 +87,13 @@ class BatterySourceType(TypedDict):
 
 
 class GasSourceType(TypedDict):
-    """Dictionary holding the source of gas storage."""
+    """Dictionary holding the source of gas consumption."""
 
     type: Literal["gas"]
 
     stat_energy_from: str
 
-    # statistic_id of costs ($) incurred from the energy meter
+    # statistic_id of costs ($) incurred from the gas meter
     # If set to None and entity_energy_price or number_energy_price are configured,
     # an EnergyCostSensor will be automatically created
     stat_cost: str | None
@@ -103,7 +103,26 @@ class GasSourceType(TypedDict):
     number_energy_price: float | None  # Price for energy ($/m³)
 
 
-SourceType = Union[GridSourceType, SolarSourceType, BatterySourceType, GasSourceType]
+class WaterSourceType(TypedDict):
+    """Dictionary holding the source of water consumption."""
+
+    type: Literal["water"]
+
+    stat_energy_from: str
+
+    # statistic_id of costs ($) incurred from the water meter
+    # If set to None and entity_energy_price or number_energy_price are configured,
+    # an EnergyCostSensor will be automatically created
+    stat_cost: str | None
+
+    # Used to generate costs if stat_cost is set to None
+    entity_energy_price: str | None  # entity_id of an entity providing price ($/m³)
+    number_energy_price: float | None  # Price for energy ($/m³)
+
+
+SourceType = Union[
+    GridSourceType, SolarSourceType, BatterySourceType, GasSourceType, WaterSourceType
+]
 
 
 class DeviceConsumption(TypedDict):
@@ -221,6 +240,15 @@ GAS_SOURCE_SCHEMA = vol.Schema(
         vol.Optional("number_energy_price"): vol.Any(vol.Coerce(float), None),
     }
 )
+WATER_SOURCE_SCHEMA = vol.Schema(
+    {
+        vol.Required("type"): "water",
+        vol.Required("stat_energy_from"): str,
+        vol.Optional("stat_cost"): vol.Any(str, None),
+        vol.Optional("entity_energy_price"): vol.Any(str, None),
+        vol.Optional("number_energy_price"): vol.Any(vol.Coerce(float), None),
+    }
+)
 
 
 def check_type_limits(value: list[SourceType]) -> list[SourceType]:
@@ -243,6 +271,7 @@ ENERGY_SOURCE_SCHEMA = vol.All(
                     "solar": SOLAR_SOURCE_SCHEMA,
                     "battery": BATTERY_SOURCE_SCHEMA,
                     "gas": GAS_SOURCE_SCHEMA,
+                    "water": WATER_SOURCE_SCHEMA,
                 },
             )
         ]

homeassistant/components/energy/sensor.py

@@ -19,6 +19,8 @@ from homeassistant.const import (
     ATTR_UNIT_OF_MEASUREMENT,
     VOLUME_CUBIC_FEET,
     VOLUME_CUBIC_METERS,
+    VOLUME_GALLONS,
+    VOLUME_LITERS,
     UnitOfEnergy,
 )
 from homeassistant.core import (
@@ -49,6 +51,12 @@ VALID_ENERGY_UNITS = [
     UnitOfEnergy.GIGA_JOULE,
 ]
 VALID_ENERGY_UNITS_GAS = [VOLUME_CUBIC_FEET, VOLUME_CUBIC_METERS] + VALID_ENERGY_UNITS
+VALID_VOLUME_UNITS_WATER = [
+    VOLUME_CUBIC_FEET,
+    VOLUME_CUBIC_METERS,
+    VOLUME_GALLONS,
+    VOLUME_LITERS,
+]
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -67,7 +75,7 @@ async def async_setup_platform(
 class SourceAdapter:
     """Adapter to allow sources and their flows to be used as sensors."""
 
-    source_type: Literal["grid", "gas"]
+    source_type: Literal["grid", "gas", "water"]
     flow_type: Literal["flow_from", "flow_to", None]
     stat_energy_key: Literal["stat_energy_from", "stat_energy_to"]
     total_money_key: Literal["stat_cost", "stat_compensation"]
@@ -100,6 +108,14 @@ SOURCE_ADAPTERS: Final = (
         "Cost",
         "cost",
     ),
+    SourceAdapter(
+        "water",
+        None,
+        "stat_energy_from",
+        "stat_cost",
+        "Cost",
+        "cost",
+    ),
 )
 
 
@@ -316,6 +332,10 @@ class EnergyCostSensor(SensorEntity):
             if energy_unit not in VALID_ENERGY_UNITS_GAS:
                 energy_unit = None
 
+        elif self._adapter.source_type == "water":
+            if energy_unit not in VALID_VOLUME_UNITS_WATER:
+                energy_unit = None
+
         if energy_unit == UnitOfEnergy.WATT_HOUR:
             energy_price /= 1000
         elif energy_unit == UnitOfEnergy.MEGA_WATT_HOUR:

homeassistant/components/energy/validate.py

@@ -13,6 +13,8 @@ from homeassistant.const import (
     STATE_UNKNOWN,
     VOLUME_CUBIC_FEET,
     VOLUME_CUBIC_METERS,
+    VOLUME_GALLONS,
+    VOLUME_LITERS,
     UnitOfEnergy,
 )
 from homeassistant.core import HomeAssistant, callback, valid_entity_id
@@ -52,6 +54,20 @@ GAS_PRICE_UNITS = tuple(
 )
 GAS_UNIT_ERROR = "entity_unexpected_unit_gas"
 GAS_PRICE_UNIT_ERROR = "entity_unexpected_unit_gas_price"
+WATER_USAGE_DEVICE_CLASSES = (sensor.SensorDeviceClass.WATER,)
+WATER_USAGE_UNITS = {
+    sensor.SensorDeviceClass.WATER: (
+        VOLUME_CUBIC_METERS,
+        VOLUME_CUBIC_FEET,
+        VOLUME_GALLONS,
+        VOLUME_LITERS,
+    ),
+}
+WATER_PRICE_UNITS = tuple(
+    f"/{unit}" for units in WATER_USAGE_UNITS.values() for unit in units
+)
+WATER_UNIT_ERROR = "entity_unexpected_unit_water"
+WATER_PRICE_UNIT_ERROR = "entity_unexpected_unit_water_price"
 
 
 @dataclasses.dataclass
@@ -437,6 +453,57 @@ async def async_validate(hass: HomeAssistant) -> EnergyPreferencesValidation:
                     )
                 )
 
+        elif source["type"] == "water":
+            wanted_statistics_metadata.add(source["stat_energy_from"])
+            validate_calls.append(
+                functools.partial(
+                    _async_validate_usage_stat,
+                    hass,
+                    statistics_metadata,
+                    source["stat_energy_from"],
+                    WATER_USAGE_DEVICE_CLASSES,
+                    WATER_USAGE_UNITS,
+                    WATER_UNIT_ERROR,
+                    source_result,
+                )
+            )
+
+            if (stat_cost := source.get("stat_cost")) is not None:
+                wanted_statistics_metadata.add(stat_cost)
+                validate_calls.append(
+                    functools.partial(
+                        _async_validate_cost_stat,
+                        hass,
+                        statistics_metadata,
+                        stat_cost,
+                        source_result,
+                    )
+                )
+            elif source.get("entity_energy_price") is not None:
+                validate_calls.append(
+                    functools.partial(
+                        _async_validate_price_entity,
+                        hass,
+                        source["entity_energy_price"],
+                        source_result,
+                        WATER_PRICE_UNITS,
+                        WATER_PRICE_UNIT_ERROR,
+                    )
+                )
+
+            if (
+                source.get("entity_energy_price") is not None
+                or source.get("number_energy_price") is not None
+            ):
+                validate_calls.append(
+                    functools.partial(
+                        _async_validate_auto_generated_cost_entity,
+                        hass,
+                        source["stat_energy_from"],
+                        source_result,
+                    )
+                )
+
         elif source["type"] == "solar":
             wanted_statistics_metadata.add(source["stat_energy_from"])
             validate_calls.append(

tests/components/energy/test_sensor.py

@@ -933,6 +933,57 @@ async def test_cost_sensor_handle_gas_kwh(
     assert state.state == "50.0"
 
 
+@pytest.mark.parametrize("unit", (VOLUME_CUBIC_FEET, VOLUME_CUBIC_METERS))
+async def test_cost_sensor_handle_water(
+    setup_integration, hass, hass_storage, unit
+) -> None:
+    """Test water cost price from sensor entity."""
+    energy_attributes = {
+        ATTR_UNIT_OF_MEASUREMENT: unit,
+        ATTR_STATE_CLASS: SensorStateClass.TOTAL_INCREASING,
+    }
+    energy_data = data.EnergyManager.default_preferences()
+    energy_data["energy_sources"].append(
+        {
+            "type": "water",
+            "stat_energy_from": "sensor.water_consumption",
+            "stat_cost": None,
+            "entity_energy_price": None,
+            "number_energy_price": 0.5,
+        }
+    )
+
+    hass_storage[data.STORAGE_KEY] = {
+        "version": 1,
+        "data": energy_data,
+    }
+
+    now = dt_util.utcnow()
+
+    hass.states.async_set(
+        "sensor.water_consumption",
+        100,
+        energy_attributes,
+    )
+
+    with patch("homeassistant.util.dt.utcnow", return_value=now):
+        await setup_integration(hass)
+
+    state = hass.states.get("sensor.water_consumption_cost")
+    assert state.state == "0.0"
+
+    # water use bumped to 200 ft³/m³
+    hass.states.async_set(
+        "sensor.water_consumption",
+        200,
+        energy_attributes,
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.water_consumption_cost")
+    assert state.state == "50.0"
+
+
 @pytest.mark.parametrize("state_class", [None])
 async def test_cost_sensor_wrong_state_class(
     setup_integration, hass, hass_storage, caplog, state_class

tests/components/energy/test_validate.py

@@ -946,3 +946,162 @@ async def test_validation_grid_no_costs_tracking(
         "energy_sources": [[]],
         "device_consumption": [],
     }
+
+
+async def test_validation_water(
+    hass, mock_energy_manager, mock_is_entity_recorded, mock_get_metadata
+):
+    """Test validating water with sensors for energy and cost/compensation."""
+    mock_is_entity_recorded["sensor.water_cost_1"] = False
+    mock_is_entity_recorded["sensor.water_compensation_1"] = False
+    await mock_energy_manager.async_update(
+        {
+            "energy_sources": [
+                {
+                    "type": "water",
+                    "stat_energy_from": "sensor.water_consumption_1",
+                    "stat_cost": "sensor.water_cost_1",
+                },
+                {
+                    "type": "water",
+                    "stat_energy_from": "sensor.water_consumption_2",
+                    "stat_cost": "sensor.water_cost_2",
+                },
+                {
+                    "type": "water",
+                    "stat_energy_from": "sensor.water_consumption_3",
+                    "stat_cost": "sensor.water_cost_2",
+                },
+                {
+                    "type": "water",
+                    "stat_energy_from": "sensor.water_consumption_4",
+                    "entity_energy_price": "sensor.water_price_1",
+                },
+                {
+                    "type": "water",
+                    "stat_energy_from": "sensor.water_consumption_3",
+                    "entity_energy_price": "sensor.water_price_2",
+                },
+            ]
+        }
+    )
+    await hass.async_block_till_done()
+    hass.states.async_set(
+        "sensor.water_consumption_1",
+        "10.10",
+        {
+            "device_class": "water",
+            "unit_of_measurement": "beers",
+            "state_class": "total_increasing",
+        },
+    )
+    hass.states.async_set(
+        "sensor.water_consumption_2",
+        "10.10",
+        {
+            "device_class": "water",
+            "unit_of_measurement": "ft³",
+            "state_class": "total_increasing",
+        },
+    )
+    hass.states.async_set(
+        "sensor.water_consumption_3",
+        "10.10",
+        {
+            "device_class": "water",
+            "unit_of_measurement": "m³",
+            "state_class": "total_increasing",
+        },
+    )
+    hass.states.async_set(
+        "sensor.water_consumption_4",
+        "10.10",
+        {"unit_of_measurement": "beers", "state_class": "total_increasing"},
+    )
+    hass.states.async_set(
+        "sensor.water_cost_2",
+        "10.10",
+        {"unit_of_measurement": "EUR/kWh", "state_class": "total_increasing"},
+    )
+    hass.states.async_set(
+        "sensor.water_price_1",
+        "10.10",
+        {"unit_of_measurement": "EUR/m³", "state_class": "total_increasing"},
+    )
+    hass.states.async_set(
+        "sensor.water_price_2",
+        "10.10",
+        {"unit_of_measurement": "EUR/invalid", "state_class": "total_increasing"},
+    )
+
+    assert (await validate.async_validate(hass)).as_dict() == {
+        "energy_sources": [
+            [
+                {
+                    "type": "entity_unexpected_unit_water",
+                    "identifier": "sensor.water_consumption_1",
+                    "value": "beers",
+                },
+                {
+                    "type": "recorder_untracked",
+                    "identifier": "sensor.water_cost_1",
+                    "value": None,
+                },
+                {
+                    "type": "entity_not_defined",
+                    "identifier": "sensor.water_cost_1",
+                    "value": None,
+                },
+            ],
+            [],
+            [],
+            [
+                {
+                    "type": "entity_unexpected_device_class",
+                    "identifier": "sensor.water_consumption_4",
+                    "value": None,
+                },
+            ],
+            [
+                {
+                    "type": "entity_unexpected_unit_water_price",
+                    "identifier": "sensor.water_price_2",
+                    "value": "EUR/invalid",
+                },
+            ],
+        ],
+        "device_consumption": [],
+    }
+
+
+async def test_validation_water_no_costs_tracking(
+    hass, mock_energy_manager, mock_is_entity_recorded, mock_get_metadata
+):
+    """Test validating water with sensors without cost tracking."""
+    await mock_energy_manager.async_update(
+        {
+            "energy_sources": [
+                {
+                    "type": "water",
+                    "stat_energy_from": "sensor.water_consumption_1",
+                    "stat_cost": None,
+                    "entity_energy_price": None,
+                    "number_energy_price": None,
+                },
+            ]
+        }
+    )
+    hass.states.async_set(
+        "sensor.water_consumption_1",
+        "10.10",
+        {
+            "device_class": "water",
+            "unit_of_measurement": "m³",
+            "state_class": "total_increasing",
+        },
+    )
+
+    assert (await validate.async_validate(hass)).as_dict() == {
+        "energy_sources": [[]],
+        "device_consumption": [],
+    }