Add monthly forecast sensor to RymPro (#101012)

* Add monthly forecast * Apply suggestions from code review Co-authored-by: epenet <[email protected]> * Remove state class and add precision * Apply suggestions from code review Co-authored-by: Robert Resch <[email protected]> * Fix ruff and mypy --------- Co-authored-by: epenet <[email protected]> Co-authored-by: Robert Resch <[email protected]>
home-assistant · Feb 15, 2024 · d4be663 · d4be663
1 parent d7787cd
commit d4be663
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 8 deletions.
diff --git a/homeassistant/components/rympro/coordinator.py b/homeassistant/components/rympro/coordinator.py
@@ -33,7 +33,12 @@ def __init__(self, hass: HomeAssistant, rympro: RymPro) -> None:
     async def _async_update_data(self) -> dict[int, dict]:
         """Fetch data from Rym Pro."""
         try:
-            return await self.rympro.last_read()
+            meters = await self.rympro.last_read()
+            for meter_id, meter in meters.items():
+                meter["consumption_forecast"] = await self.rympro.consumption_forecast(
+                    meter_id
+                )
+            return meters
         except UnauthorizedError as error:
             assert self.config_entry
             await self.hass.config_entries.async_reload(self.config_entry.entry_id)

diff --git a/homeassistant/components/rympro/sensor.py b/homeassistant/components/rympro/sensor.py
@@ -1,9 +1,12 @@
 """Sensor for RymPro meters."""
 from __future__ import annotations
 
+from dataclasses import dataclass
+
 from homeassistant.components.sensor import (
     SensorDeviceClass,
     SensorEntity,
+    SensorEntityDescription,
     SensorStateClass,
 )
 from homeassistant.config_entries import ConfigEntry
@@ -17,6 +20,30 @@
 from .coordinator import RymProDataUpdateCoordinator
 
 
+@dataclass(kw_only=True, frozen=True)
+class RymProSensorEntityDescription(SensorEntityDescription):
+    """Class describing RymPro sensor entities."""
+
+    value_key: str
+
+
+SENSOR_DESCRIPTIONS: tuple[RymProSensorEntityDescription, ...] = (
+    RymProSensorEntityDescription(
+        key="total_consumption",
+        translation_key="total_consumption",
+        state_class=SensorStateClass.TOTAL_INCREASING,
+        suggested_display_precision=3,
+        value_key="read",
+    ),
+    RymProSensorEntityDescription(
+        key="monthly_forecast",
+        translation_key="monthly_forecast",
+        suggested_display_precision=3,
+        value_key="consumption_forecast",
+    ),
+)
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
@@ -25,41 +52,41 @@ async def async_setup_entry(
     """Set up sensors for device."""
     coordinator: RymProDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
     async_add_entities(
-        RymProSensor(coordinator, meter_id, meter["read"], config_entry.entry_id)
+        RymProSensor(coordinator, meter_id, description, config_entry.entry_id)
         for meter_id, meter in coordinator.data.items()
+        for description in SENSOR_DESCRIPTIONS
     )
 
 
 class RymProSensor(CoordinatorEntity[RymProDataUpdateCoordinator], SensorEntity):
     """Sensor for RymPro meters."""
 
     _attr_has_entity_name = True
-    _attr_translation_key = "total_consumption"
     _attr_device_class = SensorDeviceClass.WATER
     _attr_native_unit_of_measurement = UnitOfVolume.CUBIC_METERS
-    _attr_state_class = SensorStateClass.TOTAL_INCREASING
+    entity_description: RymProSensorEntityDescription
 
     def __init__(
         self,
         coordinator: RymProDataUpdateCoordinator,
         meter_id: int,
-        last_read: int,
+        description: RymProSensorEntityDescription,
         entry_id: str,
     ) -> None:
         """Initialize sensor."""
         super().__init__(coordinator)
         self._meter_id = meter_id
         unique_id = f"{entry_id}_{meter_id}"
-        self._attr_unique_id = f"{unique_id}_total_consumption"
+        self._attr_unique_id = f"{unique_id}_{description.key}"
         self._attr_extra_state_attributes = {"meter_id": str(meter_id)}
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, unique_id)},
             manufacturer="Read Your Meter Pro",
             name=f"Meter {meter_id}",
         )
-        self._attr_native_value = last_read
+        self.entity_description = description
 
     @property
     def native_value(self) -> float | None:
         """Return the state of the sensor."""
-        return self.coordinator.data[self._meter_id]["read"]
+        return self.coordinator.data[self._meter_id][self.entity_description.value_key]
diff --git a/homeassistant/components/rympro/strings.json b/homeassistant/components/rympro/strings.json
@@ -21,6 +21,9 @@
     "sensor": {
       "total_consumption": {
         "name": "Total consumption"
+      },
+      "monthly_forecast": {
+        "name": "Monthly forecast"
       }
     }
   }