From 8fc8220924f318bcd49468a6e9be2e84b0272daf Mon Sep 17 00:00:00 2001 From: David Knowles Date: Mon, 28 Jul 2025 10:06:15 -0400 Subject: [PATCH] Teach Hydrawise to auto-add/remove devices (#149547) Co-authored-by: Joost Lekkerkerker --- .../components/hydrawise/binary_sensor.py | 57 ++++++---- homeassistant/components/hydrawise/const.py | 1 + .../components/hydrawise/coordinator.py | 92 ++++++++++++++- homeassistant/components/hydrawise/entity.py | 6 +- homeassistant/components/hydrawise/sensor.py | 83 +++++++++----- homeassistant/components/hydrawise/switch.py | 23 ++-- homeassistant/components/hydrawise/valve.py | 22 +++- tests/components/hydrawise/test_init.py | 106 +++++++++++++++++- 8 files changed, 321 insertions(+), 69 deletions(-) diff --git a/homeassistant/components/hydrawise/binary_sensor.py b/homeassistant/components/hydrawise/binary_sensor.py index 45537a2cc73..f2177d2144a 100644 --- a/homeassistant/components/hydrawise/binary_sensor.py +++ b/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( diff --git a/homeassistant/components/hydrawise/const.py b/homeassistant/components/hydrawise/const.py index beaf450a586..502fd14cfbd 100644 --- a/homeassistant/components/hydrawise/const.py +++ b/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) diff --git a/homeassistant/components/hydrawise/coordinator.py b/homeassistant/components/hydrawise/coordinator.py index 15d286801f9..308ffc23e36 100644 --- a/homeassistant/components/hydrawise/coordinator.py +++ b/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. diff --git a/homeassistant/components/hydrawise/entity.py b/homeassistant/components/hydrawise/entity.py index 67dd6375b0e..58153d43634 100644 --- a/homeassistant/components/hydrawise/entity.py +++ b/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, ) diff --git a/homeassistant/components/hydrawise/sensor.py b/homeassistant/components/hydrawise/sensor.py index ce0bc5a0997..3a04a587bb4 100644 --- a/homeassistant/components/hydrawise/sensor.py +++ b/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): diff --git a/homeassistant/components/hydrawise/switch.py b/homeassistant/components/hydrawise/switch.py index 7a77f27265b..238e249e1f6 100644 --- a/homeassistant/components/hydrawise/switch.py +++ b/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): diff --git a/homeassistant/components/hydrawise/valve.py b/homeassistant/components/hydrawise/valve.py index 85a91c807b2..56dd56e7d21 100644 --- a/homeassistant/components/hydrawise/valve.py +++ b/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): diff --git a/tests/components/hydrawise/test_init.py b/tests/components/hydrawise/test_init.py index 8ec3c3da648..31e86589543 100644 --- a/tests/components/hydrawise/test_init.py +++ b/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