diff --git a/homeassistant/components/husqvarna_automower/__init__.py b/homeassistant/components/husqvarna_automower/__init__.py index 0bb58fa4563..822f81f5f75 100644 --- a/homeassistant/components/husqvarna_automower/__init__.py +++ b/homeassistant/components/husqvarna_automower/__init__.py @@ -13,6 +13,7 @@ from homeassistant.helpers import ( aiohttp_client, config_entry_oauth2_flow, device_registry as dr, + entity_registry as er, ) from homeassistant.util import dt as dt_util @@ -99,3 +100,20 @@ def cleanup_removed_devices( device_reg.async_update_device( device.id, remove_config_entry_id=config_entry.entry_id ) + + +def remove_work_area_entities( + hass: HomeAssistant, + config_entry: ConfigEntry, + removed_work_areas: set[int], + mower_id: str, +) -> None: + """Remove all unused work area entities for the specified mower.""" + entity_reg = er.async_get(hass) + for entity_entry in er.async_entries_for_config_entry( + entity_reg, config_entry.entry_id + ): + for work_area_id in removed_work_areas: + if entity_entry.unique_id.startswith(f"{mower_id}_{work_area_id}_"): + _LOGGER.info("Deleting: %s", entity_entry.entity_id) + entity_reg.async_remove(entity_entry.entity_id) diff --git a/homeassistant/components/husqvarna_automower/entity.py b/homeassistant/components/husqvarna_automower/entity.py index 1bf9c004966..da6c0ae59ce 100644 --- a/homeassistant/components/husqvarna_automower/entity.py +++ b/homeassistant/components/husqvarna_automower/entity.py @@ -9,13 +9,12 @@ from typing import TYPE_CHECKING, Any from aioautomower.exceptions import ApiException from aioautomower.model import MowerActivities, MowerAttributes, MowerStates, WorkArea -from homeassistant.core import HomeAssistant, callback +from homeassistant.core import callback from homeassistant.exceptions import HomeAssistantError -from homeassistant.helpers import entity_registry as er from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import AutomowerConfigEntry, AutomowerDataUpdateCoordinator +from . import AutomowerDataUpdateCoordinator from .const import DOMAIN, EXECUTION_TIME_DELAY _LOGGER = logging.getLogger(__name__) @@ -53,30 +52,6 @@ def _work_area_translation_key(work_area_id: int, key: str) -> str: return f"work_area_{key}" -@callback -def async_remove_work_area_entities( - hass: HomeAssistant, - coordinator: AutomowerDataUpdateCoordinator, - entry: AutomowerConfigEntry, - mower_id: str, -) -> None: - """Remove deleted work areas from Home Assistant.""" - entity_reg = er.async_get(hass) - active_work_areas = set() - _work_areas = coordinator.data[mower_id].work_areas - if _work_areas is not None: - for work_area_id in _work_areas: - uid = f"{mower_id}_{work_area_id}_cutting_height_work_area" - active_work_areas.add(uid) - for entity_entry in er.async_entries_for_config_entry(entity_reg, entry.entry_id): - if ( - (split := entity_entry.unique_id.split("_"))[0] == mower_id - and split[-1] == "area" - and entity_entry.unique_id not in active_work_areas - ): - entity_reg.async_remove(entity_entry.entity_id) - - def handle_sending_exception( poll_after_sending: bool = False, ) -> Callable[ diff --git a/homeassistant/components/husqvarna_automower/number.py b/homeassistant/components/husqvarna_automower/number.py index 2a67400d1bf..d6d794f2d83 100644 --- a/homeassistant/components/husqvarna_automower/number.py +++ b/homeassistant/components/husqvarna_automower/number.py @@ -13,13 +13,12 @@ from homeassistant.const import PERCENTAGE, EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import AutomowerConfigEntry +from . import AutomowerConfigEntry, remove_work_area_entities from .coordinator import AutomowerDataUpdateCoordinator from .entity import ( AutomowerControlEntity, WorkAreaControlEntity, _work_area_translation_key, - async_remove_work_area_entities, handle_sending_exception, ) @@ -110,26 +109,44 @@ async def async_setup_entry( ) -> None: """Set up number platform.""" coordinator = entry.runtime_data - entities: list[NumberEntity] = [] + current_work_areas: dict[str, set[int]] = {} - for mower_id in coordinator.data: - if coordinator.data[mower_id].capabilities.work_areas: - _work_areas = coordinator.data[mower_id].work_areas - if _work_areas is not None: - entities.extend( - WorkAreaNumberEntity( - mower_id, coordinator, description, work_area_id + async_add_entities( + AutomowerNumberEntity(mower_id, coordinator, description) + for mower_id in coordinator.data + for description in MOWER_NUMBER_TYPES + if description.exists_fn(coordinator.data[mower_id]) + ) + + def _async_work_area_listener() -> None: + """Listen for new work areas and add/remove entities as needed.""" + for mower_id in coordinator.data: + if ( + coordinator.data[mower_id].capabilities.work_areas + and (_work_areas := coordinator.data[mower_id].work_areas) is not None + ): + received_work_areas = set(_work_areas.keys()) + current_work_area_set = current_work_areas.setdefault(mower_id, set()) + + new_work_areas = received_work_areas - current_work_area_set + removed_work_areas = current_work_area_set - received_work_areas + + if new_work_areas: + current_work_area_set.update(new_work_areas) + async_add_entities( + WorkAreaNumberEntity( + mower_id, coordinator, description, work_area_id + ) + for description in WORK_AREA_NUMBER_TYPES + for work_area_id in new_work_areas ) - for description in WORK_AREA_NUMBER_TYPES - for work_area_id in _work_areas - ) - async_remove_work_area_entities(hass, coordinator, entry, mower_id) - entities.extend( - AutomowerNumberEntity(mower_id, coordinator, description) - for description in MOWER_NUMBER_TYPES - if description.exists_fn(coordinator.data[mower_id]) - ) - async_add_entities(entities) + + if removed_work_areas: + remove_work_area_entities(hass, entry, removed_work_areas, mower_id) + current_work_area_set.difference_update(removed_work_areas) + + coordinator.async_add_listener(_async_work_area_listener) + _async_work_area_listener() class AutomowerNumberEntity(AutomowerControlEntity, NumberEntity): diff --git a/homeassistant/components/husqvarna_automower/sensor.py b/homeassistant/components/husqvarna_automower/sensor.py index 4576c4152a0..ebb68033918 100644 --- a/homeassistant/components/husqvarna_automower/sensor.py +++ b/homeassistant/components/husqvarna_automower/sensor.py @@ -431,25 +431,44 @@ async def async_setup_entry( ) -> None: """Set up sensor platform.""" coordinator = entry.runtime_data - entities: list[SensorEntity] = [] - for mower_id in coordinator.data: - if coordinator.data[mower_id].capabilities.work_areas: - _work_areas = coordinator.data[mower_id].work_areas - if _work_areas is not None: - entities.extend( - WorkAreaSensorEntity( - mower_id, coordinator, description, work_area_id - ) - for description in WORK_AREA_SENSOR_TYPES - for work_area_id in _work_areas - if description.exists_fn(_work_areas[work_area_id]) + current_work_areas: dict[str, set[int]] = {} + + async_add_entities( + AutomowerSensorEntity(mower_id, coordinator, description) + for mower_id, data in coordinator.data.items() + for description in MOWER_SENSOR_TYPES + if description.exists_fn(data) + ) + + def _async_work_area_listener() -> None: + """Listen for new work areas and add sensor entities if they did not exist. + + Listening for deletable work areas is managed in the number platform. + """ + for mower_id in coordinator.data: + if ( + coordinator.data[mower_id].capabilities.work_areas + and (_work_areas := coordinator.data[mower_id].work_areas) is not None + ): + received_work_areas = set(_work_areas.keys()) + new_work_areas = received_work_areas - current_work_areas.get( + mower_id, set() ) - entities.extend( - AutomowerSensorEntity(mower_id, coordinator, description) - for description in MOWER_SENSOR_TYPES - if description.exists_fn(coordinator.data[mower_id]) - ) - async_add_entities(entities) + if new_work_areas: + current_work_areas.setdefault(mower_id, set()).update( + new_work_areas + ) + async_add_entities( + WorkAreaSensorEntity( + mower_id, coordinator, description, work_area_id + ) + for description in WORK_AREA_SENSOR_TYPES + for work_area_id in new_work_areas + if description.exists_fn(_work_areas[work_area_id]) + ) + + coordinator.async_add_listener(_async_work_area_listener) + _async_work_area_listener() class AutomowerSensorEntity(AutomowerBaseEntity, SensorEntity): diff --git a/homeassistant/components/husqvarna_automower/switch.py b/homeassistant/components/husqvarna_automower/switch.py index c26348d875a..2bbe5c87624 100644 --- a/homeassistant/components/husqvarna_automower/switch.py +++ b/homeassistant/components/husqvarna_automower/switch.py @@ -6,8 +6,7 @@ from typing import TYPE_CHECKING, Any from aioautomower.model import MowerModes, StayOutZones, Zone from homeassistant.components.switch import SwitchEntity -from homeassistant.const import Platform -from homeassistant.core import HomeAssistant, callback +from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_registry as er from homeassistant.helpers.entity_platform import AddEntitiesCallback @@ -30,28 +29,82 @@ async def async_setup_entry( ) -> None: """Set up switch platform.""" coordinator = entry.runtime_data - entities: list[SwitchEntity] = [] - entities.extend( + current_work_areas: dict[str, set[int]] = {} + current_stay_out_zones: dict[str, set[str]] = {} + + async_add_entities( AutomowerScheduleSwitchEntity(mower_id, coordinator) for mower_id in coordinator.data ) - for mower_id in coordinator.data: - if coordinator.data[mower_id].capabilities.stay_out_zones: - _stay_out_zones = coordinator.data[mower_id].stay_out_zones - if _stay_out_zones is not None: - entities.extend( - StayOutZoneSwitchEntity(coordinator, mower_id, stay_out_zone_uid) - for stay_out_zone_uid in _stay_out_zones.zones + + def _async_work_area_listener() -> None: + """Listen for new work areas and add switch entities if they did not exist. + + Listening for deletable work areas is managed in the number platform. + """ + for mower_id in coordinator.data: + if ( + coordinator.data[mower_id].capabilities.work_areas + and (_work_areas := coordinator.data[mower_id].work_areas) is not None + ): + received_work_areas = set(_work_areas.keys()) + new_work_areas = received_work_areas - current_work_areas.get( + mower_id, set() ) - async_remove_entities(hass, coordinator, entry, mower_id) - if coordinator.data[mower_id].capabilities.work_areas: - _work_areas = coordinator.data[mower_id].work_areas - if _work_areas is not None: - entities.extend( - WorkAreaSwitchEntity(coordinator, mower_id, work_area_id) - for work_area_id in _work_areas + if new_work_areas: + current_work_areas.setdefault(mower_id, set()).update( + new_work_areas + ) + async_add_entities( + WorkAreaSwitchEntity(coordinator, mower_id, work_area_id) + for work_area_id in new_work_areas + ) + + def _remove_stay_out_zone_entities( + removed_stay_out_zones: set, mower_id: str + ) -> None: + """Remove all unused stay-out zones for all platforms.""" + entity_reg = er.async_get(hass) + for entity_entry in er.async_entries_for_config_entry( + entity_reg, entry.entry_id + ): + for stay_out_zone_uid in removed_stay_out_zones: + if entity_entry.unique_id.startswith(f"{mower_id}_{stay_out_zone_uid}"): + entity_reg.async_remove(entity_entry.entity_id) + + def _async_stay_out_zone_listener() -> None: + """Listen for new stay-out zones and add/remove switch entities if they did not exist.""" + for mower_id in coordinator.data: + if ( + coordinator.data[mower_id].capabilities.stay_out_zones + and (_stay_out_zones := coordinator.data[mower_id].stay_out_zones) + is not None + ): + received_stay_out_zones = set(_stay_out_zones.zones) + current_stay_out_zones_set = current_stay_out_zones.get(mower_id, set()) + new_stay_out_zones = ( + received_stay_out_zones - current_stay_out_zones_set ) - async_add_entities(entities) + removed_stay_out_zones = ( + current_stay_out_zones_set - received_stay_out_zones + ) + if new_stay_out_zones: + current_stay_out_zones.setdefault(mower_id, set()).update( + new_stay_out_zones + ) + async_add_entities( + StayOutZoneSwitchEntity( + coordinator, mower_id, stay_out_zone_uid + ) + for stay_out_zone_uid in new_stay_out_zones + ) + if removed_stay_out_zones: + _remove_stay_out_zone_entities(removed_stay_out_zones, mower_id) + + coordinator.async_add_listener(_async_work_area_listener) + coordinator.async_add_listener(_async_stay_out_zone_listener) + _async_work_area_listener() + _async_stay_out_zone_listener() class AutomowerScheduleSwitchEntity(AutomowerControlEntity, SwitchEntity): @@ -180,28 +233,3 @@ class WorkAreaSwitchEntity(WorkAreaControlEntity, SwitchEntity): await self.coordinator.api.commands.workarea_settings( self.mower_id, self.work_area_id, enabled=True ) - - -@callback -def async_remove_entities( - hass: HomeAssistant, - coordinator: AutomowerDataUpdateCoordinator, - entry: AutomowerConfigEntry, - mower_id: str, -) -> None: - """Remove deleted stay-out-zones from Home Assistant.""" - entity_reg = er.async_get(hass) - active_zones = set() - _zones = coordinator.data[mower_id].stay_out_zones - if _zones is not None: - for zones_uid in _zones.zones: - uid = f"{mower_id}_{zones_uid}_stay_out_zones" - active_zones.add(uid) - for entity_entry in er.async_entries_for_config_entry(entity_reg, entry.entry_id): - if ( - entity_entry.domain == Platform.SWITCH - and (split := entity_entry.unique_id.split("_"))[0] == mower_id - and split[-1] == "zones" - and entity_entry.unique_id not in active_zones - ): - entity_reg.async_remove(entity_entry.entity_id) diff --git a/tests/components/husqvarna_automower/test_init.py b/tests/components/husqvarna_automower/test_init.py index daebb743c2f..b2127145372 100644 --- a/tests/components/husqvarna_automower/test_init.py +++ b/tests/components/husqvarna_automower/test_init.py @@ -1,6 +1,6 @@ """Tests for init module.""" -from datetime import timedelta +from datetime import datetime, timedelta import http import time from unittest.mock import AsyncMock @@ -10,15 +10,17 @@ from aioautomower.exceptions import ( AuthException, HusqvarnaWSServerHandshakeError, ) -from aioautomower.model import MowerAttributes +from aioautomower.model import MowerAttributes, WorkArea from freezegun.api import FrozenDateTimeFactory import pytest from syrupy.assertion import SnapshotAssertion from homeassistant.components.husqvarna_automower.const import DOMAIN, OAUTH2_TOKEN +from homeassistant.components.husqvarna_automower.coordinator import SCAN_INTERVAL from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant from homeassistant.helpers import device_registry as dr, entity_registry as er +from homeassistant.util import dt as dt_util from . import setup_integration from .const import TEST_MOWER_ID @@ -26,6 +28,10 @@ from .const import TEST_MOWER_ID from tests.common import MockConfigEntry, async_fire_time_changed from tests.test_util.aiohttp import AiohttpClientMocker +ADDITIONAL_NUMBER_ENTITIES = 1 +ADDITIONAL_SENSOR_ENTITIES = 2 +ADDITIONAL_SWITCH_ENTITIES = 1 + async def test_load_unload_entry( hass: HomeAssistant, @@ -163,29 +169,6 @@ async def test_device_info( assert reg_device == snapshot -async def test_workarea_deleted( - hass: HomeAssistant, - mock_automower_client: AsyncMock, - mock_config_entry: MockConfigEntry, - entity_registry: er.EntityRegistry, - values: dict[str, MowerAttributes], -) -> None: - """Test if work area is deleted after removed.""" - - await setup_integration(hass, mock_config_entry) - current_entries = len( - er.async_entries_for_config_entry(entity_registry, mock_config_entry.entry_id) - ) - - del values[TEST_MOWER_ID].work_areas[123456] - mock_automower_client.get_status.return_value = values - await hass.config_entries.async_reload(mock_config_entry.entry_id) - await hass.async_block_till_done() - assert len( - er.async_entries_for_config_entry(entity_registry, mock_config_entry.entry_id) - ) == (current_entries - 2) - - async def test_coordinator_automatic_registry_cleanup( hass: HomeAssistant, mock_automower_client: AsyncMock, @@ -219,3 +202,70 @@ async def test_coordinator_automatic_registry_cleanup( len(dr.async_entries_for_config_entry(device_registry, entry.entry_id)) == current_devices - 1 ) + + +async def test_add_and_remove_work_area( + hass: HomeAssistant, + mock_automower_client: AsyncMock, + mock_config_entry: MockConfigEntry, + freezer: FrozenDateTimeFactory, + entity_registry: er.EntityRegistry, + values: dict[str, MowerAttributes], +) -> None: + """Test adding a work area in runtime.""" + await setup_integration(hass, mock_config_entry) + entry = hass.config_entries.async_entries(DOMAIN)[0] + current_entites_start = len( + er.async_entries_for_config_entry(entity_registry, entry.entry_id) + ) + values[TEST_MOWER_ID].work_area_names.append("new work area") + values[TEST_MOWER_ID].work_area_dict.update({1: "new work area"}) + values[TEST_MOWER_ID].work_areas.update( + { + 1: WorkArea( + name="new work area", + cutting_height=12, + enabled=True, + progress=12, + last_time_completed=datetime( + 2024, 10, 1, 11, 11, 0, tzinfo=dt_util.get_default_time_zone() + ), + ) + } + ) + mock_automower_client.get_status.return_value = values + freezer.tick(SCAN_INTERVAL) + async_fire_time_changed(hass) + await hass.async_block_till_done() + current_entites_after_addition = len( + er.async_entries_for_config_entry(entity_registry, entry.entry_id) + ) + assert ( + current_entites_after_addition + == current_entites_start + + ADDITIONAL_NUMBER_ENTITIES + + ADDITIONAL_SENSOR_ENTITIES + + ADDITIONAL_SWITCH_ENTITIES + ) + + values[TEST_MOWER_ID].work_area_names.remove("new work area") + del values[TEST_MOWER_ID].work_area_dict[1] + del values[TEST_MOWER_ID].work_areas[1] + values[TEST_MOWER_ID].work_area_names.remove("Front lawn") + del values[TEST_MOWER_ID].work_area_dict[123456] + del values[TEST_MOWER_ID].work_areas[123456] + del values[TEST_MOWER_ID].calendar.tasks[:2] + mock_automower_client.get_status.return_value = values + freezer.tick(SCAN_INTERVAL) + async_fire_time_changed(hass) + await hass.async_block_till_done() + current_entites_after_deletion = len( + er.async_entries_for_config_entry(entity_registry, entry.entry_id) + ) + assert ( + current_entites_after_deletion + == current_entites_start + - ADDITIONAL_SWITCH_ENTITIES + - ADDITIONAL_NUMBER_ENTITIES + - ADDITIONAL_SENSOR_ENTITIES + ) diff --git a/tests/components/husqvarna_automower/test_switch.py b/tests/components/husqvarna_automower/test_switch.py index 0dd5acfaf6b..100fd9fe3a4 100644 --- a/tests/components/husqvarna_automower/test_switch.py +++ b/tests/components/husqvarna_automower/test_switch.py @@ -5,7 +5,7 @@ from unittest.mock import AsyncMock, patch import zoneinfo from aioautomower.exceptions import ApiException -from aioautomower.model import MowerAttributes, MowerModes +from aioautomower.model import MowerAttributes, MowerModes, Zone from aioautomower.utils import mower_list_to_dictionary_dataclass from freezegun.api import FrozenDateTimeFactory import pytest @@ -38,8 +38,9 @@ from tests.common import ( snapshot_platform, ) -TEST_ZONE_ID = "AAAAAAAA-BBBB-CCCC-DDDD-123456789101" TEST_AREA_ID = 0 +TEST_VARIABLE_ZONE_ID = "203F6359-AB56-4D57-A6DC-703095BB695D" +TEST_ZONE_ID = "AAAAAAAA-BBBB-CCCC-DDDD-123456789101" async def test_switch_states( @@ -179,6 +180,7 @@ async def test_work_area_switch_commands( mock_config_entry: MockConfigEntry, freezer: FrozenDateTimeFactory, mower_time_zone: zoneinfo.ZoneInfo, + values: dict[str, MowerAttributes], ) -> None: """Test switch commands.""" entity_id = "switch.test_mower_1_my_lawn" @@ -219,26 +221,46 @@ async def test_work_area_switch_commands( assert len(mocked_method.mock_calls) == 2 -async def test_zones_deleted( +async def test_add_stay_out_zone( hass: HomeAssistant, mock_automower_client: AsyncMock, mock_config_entry: MockConfigEntry, + freezer: FrozenDateTimeFactory, entity_registry: er.EntityRegistry, values: dict[str, MowerAttributes], ) -> None: - """Test if stay-out-zone is deleted after removed.""" + """Test adding a stay out zone in runtime.""" await setup_integration(hass, mock_config_entry) - current_entries = len( - er.async_entries_for_config_entry(entity_registry, mock_config_entry.entry_id) + entry = hass.config_entries.async_entries(DOMAIN)[0] + current_entites = len( + er.async_entries_for_config_entry(entity_registry, entry.entry_id) + ) + values[TEST_MOWER_ID].stay_out_zones.zones.update( + { + TEST_VARIABLE_ZONE_ID: Zone( + name="future_zone", + enabled=True, + ) + } ) - - del values[TEST_MOWER_ID].stay_out_zones.zones[TEST_ZONE_ID] mock_automower_client.get_status.return_value = values - await hass.config_entries.async_reload(mock_config_entry.entry_id) + freezer.tick(SCAN_INTERVAL) + async_fire_time_changed(hass) await hass.async_block_till_done() - assert len( - er.async_entries_for_config_entry(entity_registry, mock_config_entry.entry_id) - ) == (current_entries - 1) + current_entites_after_addition = len( + er.async_entries_for_config_entry(entity_registry, entry.entry_id) + ) + assert current_entites_after_addition == current_entites + 1 + values[TEST_MOWER_ID].stay_out_zones.zones.pop(TEST_VARIABLE_ZONE_ID) + values[TEST_MOWER_ID].stay_out_zones.zones.pop(TEST_ZONE_ID) + mock_automower_client.get_status.return_value = values + freezer.tick(SCAN_INTERVAL) + async_fire_time_changed(hass) + await hass.async_block_till_done() + current_entites_after_deletion = len( + er.async_entries_for_config_entry(entity_registry, entry.entry_id) + ) + assert current_entites_after_deletion == current_entites - 1 async def test_switch_snapshot(