evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add and remove entities during runtime in Husqvarna Automower (#127878)

Browse source
This commit is contained in:
Thomas55555 2024-10-29 12:46:04 +01:00 • committed by GitHub
parent 2236ca3e12
commit 983cd9c3fc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 275 additions and 146 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
homeassistant/components/husqvarna_automower/__init__.py
View file

@ -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)

29
homeassistant/components/husqvarna_automower/entity.py
View file

@ -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[

57
homeassistant/components/husqvarna_automower/number.py
View file

@ -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):

55
homeassistant/components/husqvarna_automower/sensor.py
View file

@ -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):

116
homeassistant/components/husqvarna_automower/switch.py
View file

@ -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)

100
tests/components/husqvarna_automower/test_init.py
View file

@ -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
)

46
tests/components/husqvarna_automower/test_switch.py
View file

@ -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(