Teach Hydrawise to auto-add/remove devices (#149547)

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

homeassistant/components/hydrawise/binary_sensor.py

@@ -2,11 +2,11 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from dataclasses import dataclass
 from datetime import datetime
 
-from pydrawise import Zone
+from pydrawise import Controller, Zone
 import voluptuous as vol
 
 from homeassistant.components.binary_sensor import (
@@ -81,31 +81,46 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Hydrawise binary_sensor platform."""
     coordinators = config_entry.runtime_data
-    entities: list[HydrawiseBinarySensor] = []
-    for controller in coordinators.main.data.controllers.values():
-        entities.extend(
-            HydrawiseBinarySensor(coordinators.main, description, controller)
-            for description in CONTROLLER_BINARY_SENSORS
-        )
-        entities.extend(
-            HydrawiseBinarySensor(
-                coordinators.main,
-                description,
-                controller,
-                sensor_id=sensor.id,
+
+    def _add_new_controllers(controllers: Iterable[Controller]) -> None:
+        entities: list[HydrawiseBinarySensor] = []
+        for controller in controllers:
+            entities.extend(
+                HydrawiseBinarySensor(coordinators.main, description, controller)
+                for description in CONTROLLER_BINARY_SENSORS
             )
-            for sensor in controller.sensors
-            for description in RAIN_SENSOR_BINARY_SENSOR
-            if "rain sensor" in sensor.model.name.lower()
-        )
-        entities.extend(
+            entities.extend(
+                HydrawiseBinarySensor(
+                    coordinators.main,
+                    description,
+                    controller,
+                    sensor_id=sensor.id,
+                )
+                for sensor in controller.sensors
+                for description in RAIN_SENSOR_BINARY_SENSOR
+                if "rain sensor" in sensor.model.name.lower()
+            )
+        async_add_entities(entities)
+
+    def _add_new_zones(zones: Iterable[tuple[Zone, Controller]]) -> None:
+        async_add_entities(
             HydrawiseZoneBinarySensor(
                 coordinators.main, description, controller, zone_id=zone.id
             )
-            for zone in controller.zones
+            for zone, controller in zones
             for description in ZONE_BINARY_SENSORS
         )
-    async_add_entities(entities)
+
+    _add_new_controllers(coordinators.main.data.controllers.values())
+    _add_new_zones(
+        [
+            (zone, coordinators.main.data.zone_id_to_controller[zone.id])
+            for zone in coordinators.main.data.zones.values()
+        ]
+    )
+    coordinators.main.new_controllers_callbacks.append(_add_new_controllers)
+    coordinators.main.new_zones_callbacks.append(_add_new_zones)
+
     platform = entity_platform.async_get_current_platform()
     platform.async_register_entity_service(SERVICE_RESUME, None, "resume")
     platform.async_register_entity_service(

homeassistant/components/hydrawise/const.py

@@ -13,6 +13,7 @@ DOMAIN = "hydrawise"
 DEFAULT_WATERING_TIME = timedelta(minutes=15)
 
 MANUFACTURER = "Hydrawise"
+MODEL_ZONE = "Zone"
 
 MAIN_SCAN_INTERVAL = timedelta(minutes=5)
 WATER_USE_SCAN_INTERVAL = timedelta(minutes=60)

homeassistant/components/hydrawise/coordinator.py

@@ -2,17 +2,26 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable, Iterable
 from dataclasses import dataclass, field
 
 from pydrawise import HydrawiseBase
 from pydrawise.schema import Controller, ControllerWaterUseSummary, Sensor, User, Zone
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers.device_registry import DeviceEntry
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 from homeassistant.util.dt import now
 
-from .const import DOMAIN, LOGGER, MAIN_SCAN_INTERVAL, WATER_USE_SCAN_INTERVAL
+from .const import (
+    DOMAIN,
+    LOGGER,
+    MAIN_SCAN_INTERVAL,
+    MODEL_ZONE,
+    WATER_USE_SCAN_INTERVAL,
+)
 
 type HydrawiseConfigEntry = ConfigEntry[HydrawiseUpdateCoordinators]
 
@@ -24,6 +33,7 @@ class HydrawiseData:
     user: User
     controllers: dict[int, Controller] = field(default_factory=dict)
     zones: dict[int, Zone] = field(default_factory=dict)
+    zone_id_to_controller: dict[int, Controller] = field(default_factory=dict)
     sensors: dict[int, Sensor] = field(default_factory=dict)
     daily_water_summary: dict[int, ControllerWaterUseSummary] = field(
         default_factory=dict
@@ -68,6 +78,13 @@ class HydrawiseMainDataUpdateCoordinator(HydrawiseDataUpdateCoordinator):
             update_interval=MAIN_SCAN_INTERVAL,
         )
         self.api = api
+        self.new_controllers_callbacks: list[
+            Callable[[Iterable[Controller]], None]
+        ] = []
+        self.new_zones_callbacks: list[
+            Callable[[Iterable[tuple[Zone, Controller]]], None]
+        ] = []
+        self.async_add_listener(self._add_remove_zones)
 
     async def _async_update_data(self) -> HydrawiseData:
         """Fetch the latest data from Hydrawise."""
@@ -80,10 +97,81 @@ class HydrawiseMainDataUpdateCoordinator(HydrawiseDataUpdateCoordinator):
             controller.zones = await self.api.get_zones(controller)
             for zone in controller.zones:
                 data.zones[zone.id] = zone
+                data.zone_id_to_controller[zone.id] = controller
             for sensor in controller.sensors:
                 data.sensors[sensor.id] = sensor
         return data
 
+    @callback
+    def _add_remove_zones(self) -> None:
+        """Add newly discovered zones and remove nonexistent ones."""
+        if self.data is None:
+            # Likely a setup error; ignore.
+            # Despite what mypy thinks, this is still reachable. Without this check,
+            # the test_connect_retry test in test_init.py fails.
+            return  # type: ignore[unreachable]
+
+        device_registry = dr.async_get(self.hass)
+        devices = dr.async_entries_for_config_entry(
+            device_registry, self.config_entry.entry_id
+        )
+        previous_zones: set[str] = set()
+        previous_zones_by_id: dict[str, DeviceEntry] = {}
+        previous_controllers: set[str] = set()
+        previous_controllers_by_id: dict[str, DeviceEntry] = {}
+        for device in devices:
+            for domain, identifier in device.identifiers:
+                if domain == DOMAIN:
+                    if device.model == MODEL_ZONE:
+                        previous_zones.add(identifier)
+                        previous_zones_by_id[identifier] = device
+                    else:
+                        previous_controllers.add(identifier)
+                        previous_controllers_by_id[identifier] = device
+                    continue
+
+        current_zones = {str(zone_id) for zone_id in self.data.zones}
+        current_controllers = {
+            str(controller_id) for controller_id in self.data.controllers
+        }
+
+        if removed_zones := previous_zones - current_zones:
+            LOGGER.debug("Removed zones: %s", ", ".join(removed_zones))
+            for zone_id in removed_zones:
+                device_registry.async_update_device(
+                    device_id=previous_zones_by_id[zone_id].id,
+                    remove_config_entry_id=self.config_entry.entry_id,
+                )
+
+        if removed_controllers := previous_controllers - current_controllers:
+            LOGGER.debug("Removed controllers: %s", ", ".join(removed_controllers))
+            for controller_id in removed_controllers:
+                device_registry.async_update_device(
+                    device_id=previous_controllers_by_id[controller_id].id,
+                    remove_config_entry_id=self.config_entry.entry_id,
+                )
+
+        if new_controller_ids := current_controllers - previous_controllers:
+            LOGGER.debug("New controllers found: %s", ", ".join(new_controller_ids))
+            new_controllers = [
+                self.data.controllers[controller_id]
+                for controller_id in map(int, new_controller_ids)
+            ]
+            for new_controller_callback in self.new_controllers_callbacks:
+                new_controller_callback(new_controllers)
+
+        if new_zone_ids := current_zones - previous_zones:
+            LOGGER.debug("New zones found: %s", ", ".join(new_zone_ids))
+            new_zones = [
+                (
+                    self.data.zones[zone_id],
+                    self.data.zone_id_to_controller[zone_id],
+                )
+                for zone_id in map(int, new_zone_ids)
+            ]
+            for new_zone_callback in self.new_zones_callbacks:
+                new_zone_callback(new_zones)
+
 
 class HydrawiseWaterUseDataUpdateCoordinator(HydrawiseDataUpdateCoordinator):
     """Data Update Coordinator for Hydrawise Water Use.

homeassistant/components/hydrawise/entity.py

@@ -9,7 +9,7 @@ from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import EntityDescription
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from .const import DOMAIN, MANUFACTURER
+from .const import DOMAIN, MANUFACTURER, MODEL_ZONE
 from .coordinator import HydrawiseDataUpdateCoordinator
 
 
@@ -40,7 +40,9 @@ class HydrawiseEntity(CoordinatorEntity[HydrawiseDataUpdateCoordinator]):
             identifiers={(DOMAIN, self._device_id)},
             name=self.zone.name if zone_id is not None else controller.name,
             model=(
-                "Zone" if zone_id is not None else controller.hardware.model.description
+                MODEL_ZONE
+                if zone_id is not None
+                else controller.hardware.model.description
             ),
             manufacturer=MANUFACTURER,
         )

homeassistant/components/hydrawise/sensor.py

@@ -2,12 +2,12 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from dataclasses import dataclass
 from datetime import timedelta
 from typing import Any
 
-from pydrawise.schema import ControllerWaterUseSummary
+from pydrawise.schema import Controller, ControllerWaterUseSummary, Zone
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
@@ -31,7 +31,9 @@ class HydrawiseSensorEntityDescription(SensorEntityDescription):
 
 
 def _get_water_use(sensor: HydrawiseSensor) -> ControllerWaterUseSummary:
-    return sensor.coordinator.data.daily_water_summary[sensor.controller.id]
+    return sensor.coordinator.data.daily_water_summary.get(
+        sensor.controller.id, ControllerWaterUseSummary()
+    )
 
 
 WATER_USE_CONTROLLER_SENSORS: tuple[HydrawiseSensorEntityDescription, ...] = (
@@ -133,44 +135,65 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Hydrawise sensor platform."""
     coordinators = config_entry.runtime_data
-    entities: list[HydrawiseSensor] = []
-    for controller in coordinators.main.data.controllers.values():
-        entities.extend(
-            HydrawiseSensor(coordinators.water_use, description, controller)
-            for description in WATER_USE_CONTROLLER_SENSORS
+
+    def _has_flow_sensor(controller: Controller) -> bool:
+        daily_water_use_summary = coordinators.water_use.data.daily_water_summary.get(
+            controller.id, ControllerWaterUseSummary()
         )
-        entities.extend(
-            HydrawiseSensor(
-                coordinators.water_use, description, controller, zone_id=zone.id
-            )
-            for zone in controller.zones
-            for description in WATER_USE_ZONE_SENSORS
-        )
-        entities.extend(
-            HydrawiseSensor(coordinators.main, description, controller, zone_id=zone.id)
-            for zone in controller.zones
-            for description in ZONE_SENSORS
-        )
-        if (
-            coordinators.water_use.data.daily_water_summary[controller.id].total_use
-            is not None
-        ):
-            # we have a flow sensor for this controller
+        return daily_water_use_summary.total_use is not None
+
+    def _add_new_controllers(controllers: Iterable[Controller]) -> None:
+        entities: list[HydrawiseSensor] = []
+        for controller in controllers:
             entities.extend(
                 HydrawiseSensor(coordinators.water_use, description, controller)
-                for description in FLOW_CONTROLLER_SENSORS
+                for description in WATER_USE_CONTROLLER_SENSORS
             )
-            entities.extend(
+            if _has_flow_sensor(controller):
+                entities.extend(
+                    HydrawiseSensor(coordinators.water_use, description, controller)
+                    for description in FLOW_CONTROLLER_SENSORS
+                )
+        async_add_entities(entities)
+
+    def _add_new_zones(zones: Iterable[tuple[Zone, Controller]]) -> None:
+        async_add_entities(
+            [
+                HydrawiseSensor(
+                    coordinators.water_use, description, controller, zone_id=zone.id
+                )
+                for zone, controller in zones
+                for description in WATER_USE_ZONE_SENSORS
+            ]
+            + [
+                HydrawiseSensor(
+                    coordinators.main, description, controller, zone_id=zone.id
+                )
+                for zone, controller in zones
+                for description in ZONE_SENSORS
+            ]
+            + [
                 HydrawiseSensor(
                     coordinators.water_use,
                     description,
                     controller,
                     zone_id=zone.id,
                 )
-                for zone in controller.zones
+                for zone, controller in zones
                 for description in FLOW_ZONE_SENSORS
-            )
-    async_add_entities(entities)
+                if _has_flow_sensor(controller)
+            ]
+        )
+
+    _add_new_controllers(coordinators.main.data.controllers.values())
+    _add_new_zones(
+        [
+            (zone, coordinators.main.data.zone_id_to_controller[zone.id])
+            for zone in coordinators.main.data.zones.values()
+        ]
+    )
+    coordinators.main.new_controllers_callbacks.append(_add_new_controllers)
+    coordinators.main.new_zones_callbacks.append(_add_new_zones)
 
 
 class HydrawiseSensor(HydrawiseEntity, SensorEntity):

homeassistant/components/hydrawise/switch.py

@@ -2,12 +2,12 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Coroutine
+from collections.abc import Callable, Coroutine, Iterable
 from dataclasses import dataclass
 from datetime import timedelta
 from typing import Any
 
-from pydrawise import HydrawiseBase, Zone
+from pydrawise import Controller, HydrawiseBase, Zone
 
 from homeassistant.components.switch import (
     SwitchDeviceClass,
@@ -66,12 +66,21 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Hydrawise switch platform."""
     coordinators = config_entry.runtime_data
-    async_add_entities(
-        HydrawiseSwitch(coordinators.main, description, controller, zone_id=zone.id)
-        for controller in coordinators.main.data.controllers.values()
-        for zone in controller.zones
-        for description in SWITCH_TYPES
+
+    def _add_new_zones(zones: Iterable[tuple[Zone, Controller]]) -> None:
+        async_add_entities(
+            HydrawiseSwitch(coordinators.main, description, controller, zone_id=zone.id)
+            for zone, controller in zones
+            for description in SWITCH_TYPES
+        )
+
+    _add_new_zones(
+        [
+            (zone, coordinators.main.data.zone_id_to_controller[zone.id])
+            for zone in coordinators.main.data.zones.values()
+        ]
     )
+    coordinators.main.new_zones_callbacks.append(_add_new_zones)
 
 
 class HydrawiseSwitch(HydrawiseEntity, SwitchEntity):

homeassistant/components/hydrawise/valve.py

@@ -2,9 +2,10 @@
 
 from __future__ import annotations
 
+from collections.abc import Iterable
 from typing import Any
 
-from pydrawise.schema import Zone
+from pydrawise.schema import Controller, Zone
 
 from homeassistant.components.valve import (
     ValveDeviceClass,
@@ -33,12 +34,21 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Hydrawise valve platform."""
     coordinators = config_entry.runtime_data
-    async_add_entities(
-        HydrawiseValve(coordinators.main, description, controller, zone_id=zone.id)
-        for controller in coordinators.main.data.controllers.values()
-        for zone in controller.zones
-        for description in VALVE_TYPES
+
+    def _add_new_zones(zones: Iterable[tuple[Zone, Controller]]) -> None:
+        async_add_entities(
+            HydrawiseValve(coordinators.main, description, controller, zone_id=zone.id)
+            for zone, controller in zones
+            for description in VALVE_TYPES
+        )
+
+    _add_new_zones(
+        [
+            (zone, coordinators.main.data.zone_id_to_controller[zone.id])
+            for zone in coordinators.main.data.zones.values()
+        ]
     )
+    coordinators.main.new_zones_callbacks.append(_add_new_zones)
 
 
 class HydrawiseValve(HydrawiseEntity, ValveEntity):

tests/components/hydrawise/test_init.py

@@ -1,13 +1,19 @@
 """Tests for the Hydrawise integration."""
 
+from copy import deepcopy
 from unittest.mock import AsyncMock
 
 from aiohttp import ClientError
+from freezegun.api import FrozenDateTimeFactory
+from pydrawise.schema import Controller, User, Zone
 
+from homeassistant.components.hydrawise.const import DOMAIN, MAIN_SCAN_INTERVAL
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers.device_registry import DeviceRegistry
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 
 async def test_connect_retry(
@@ -32,3 +38,101 @@ async def test_update_version(
 
     # Make sure reauth flow has been initiated
     assert any(mock_config_entry_legacy.async_get_active_flows(hass, {"reauth"}))
+
+
+async def test_auto_add_devices(
+    hass: HomeAssistant,
+    device_registry: DeviceRegistry,
+    mock_added_config_entry: MockConfigEntry,
+    mock_pydrawise: AsyncMock,
+    user: User,
+    controller: Controller,
+    zones: list[Zone],
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test new devices are auto-added to the device registry."""
+    device = device_registry.async_get_device(
+        identifiers={(DOMAIN, str(controller.id))}
+    )
+    assert device is not None
+    for zone in zones:
+        zone_device = device_registry.async_get_device(
+            identifiers={(DOMAIN, str(zone.id))}
+        )
+        assert zone_device is not None
+    all_devices = dr.async_entries_for_config_entry(
+        device_registry, mock_added_config_entry.entry_id
+    )
+    # 1 controller + 2 zones
+    assert len(all_devices) == 3
+
+    controller2 = deepcopy(controller)
+    controller2.id += 10
+    controller2.name += " 2"
+    controller2.sensors = []
+
+    zones2 = deepcopy(zones)
+    for zone in zones2:
+        zone.id += 10
+        zone.name += " 2"
+
+    user.controllers = [controller, controller2]
+    mock_pydrawise.get_zones.side_effect = [zones, zones2]
+
+    # Make the coordinator refresh data.
+    freezer.tick(MAIN_SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    new_controller_device = device_registry.async_get_device(
+        identifiers={(DOMAIN, str(controller2.id))}
+    )
+    assert new_controller_device is not None
+    for zone in zones2:
+        new_zone_device = device_registry.async_get_device(
+            identifiers={(DOMAIN, str(zone.id))}
+        )
+        assert new_zone_device is not None
+
+    all_devices = dr.async_entries_for_config_entry(
+        device_registry, mock_added_config_entry.entry_id
+    )
+    # 2 controllers + 4 zones
+    assert len(all_devices) == 6
+
+
+async def test_auto_remove_devices(
+    hass: HomeAssistant,
+    device_registry: DeviceRegistry,
+    mock_added_config_entry: MockConfigEntry,
+    user: User,
+    controller: Controller,
+    zones: list[Zone],
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test old devices are auto-removed from the device registry."""
+    assert (
+        device_registry.async_get_device(identifiers={(DOMAIN, str(controller.id))})
+        is not None
+    )
+    for zone in zones:
+        device = device_registry.async_get_device(identifiers={(DOMAIN, str(zone.id))})
+        assert device is not None
+
+    user.controllers = []
+    # Make the coordinator refresh data.
+    freezer.tick(MAIN_SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    assert (
+        device_registry.async_get_device(identifiers={(DOMAIN, str(controller.id))})
+        is None
+    )
+    for zone in zones:
+        device = device_registry.async_get_device(identifiers={(DOMAIN, str(zone.id))})
+        assert device is None
+    all_devices = dr.async_entries_for_config_entry(
+        device_registry, mock_added_config_entry.entry_id
+    )
+    assert len(all_devices) == 0