Skip to content

Commit

Permalink
Add work area switch for Husqvarna Automower (#126376)
Browse files Browse the repository at this point in the history 
* Add work area switch for Husqvarna Automower

* move work area deletion test to separate file

* stale doctsrings

* don't use custom test file

* use _attr_name

* ruff

* add available property

* hassfest

* fix tests

* Apply suggestions from code review

Co-authored-by: Joost Lekkerkerker <[email protected]>

* constants

---------

Co-authored-by: Joost Lekkerkerker <[email protected]>
  • Loading branch information
@Thomas55555 @joostlek
Thomas55555 and joostlek authored Sep 24, 2024
1 parent c9351fd commit dc77b2d
Show file tree
Hide file tree
Showing 9 changed files with 404 additions and 105 deletions.
72 changes: 68 additions & 4 deletions homeassistant/components/husqvarna_automower/entity.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,18 @@
from collections.abc import Awaitable, Callable, Coroutine
import functools
import logging
from typing import Any
from typing import TYPE_CHECKING, Any

from aioautomower.exceptions import ApiException
from aioautomower.model import MowerActivities, MowerAttributes, MowerStates
from aioautomower.model import MowerActivities, MowerAttributes, MowerStates, WorkArea

from homeassistant.core import callback
from homeassistant.core import HomeAssistant, 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 AutomowerDataUpdateCoordinator
from . import AutomowerConfigEntry, AutomowerDataUpdateCoordinator
from .const import DOMAIN, EXECUTION_TIME_DELAY

_LOGGER = logging.getLogger(__name__)
Expand Down Expand Up @@ -44,6 +45,38 @@ def _check_error_free(mower_attributes: MowerAttributes) -> bool:
)


@callback
def _work_area_translation_key(work_area_id: int, key: str) -> str:
"""Return the translation key."""
if work_area_id == 0:
return f"my_lawn_{key}"
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[
Expand Down Expand Up @@ -120,3 +153,34 @@ class AutomowerControlEntity(AutomowerAvailableEntity):
def available(self) -> bool:
"""Return True if the device is available."""
return super().available and _check_error_free(self.mower_attributes)


class WorkAreaControlEntity(AutomowerControlEntity):
"""Base entity work work areas with control function."""

def __init__(
self,
mower_id: str,
coordinator: AutomowerDataUpdateCoordinator,
work_area_id: int,
) -> None:
"""Initialize AutomowerEntity."""
super().__init__(mower_id, coordinator)
self.work_area_id = work_area_id

@property
def work_areas(self) -> dict[int, WorkArea]:
"""Get the work areas from the mower attributes."""
if TYPE_CHECKING:
assert self.mower_attributes.work_areas is not None
return self.mower_attributes.work_areas

@property
def work_area_attributes(self) -> WorkArea:
"""Get the work area attributes of the current work area."""
return self.work_areas[self.work_area_id]

@property
def available(self) -> bool:
"""Return True if the work area is available and the mower has no errors."""
return super().available and self.work_area_id in self.work_areas
70 changes: 19 additions & 51 deletions homeassistant/components/husqvarna_automower/number.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,14 +9,19 @@
from aioautomower.session import AutomowerSession

from homeassistant.components.number import NumberEntity, NumberEntityDescription
from homeassistant.const import PERCENTAGE, EntityCategory, Platform
from homeassistant.const import PERCENTAGE, EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.entity_platform import AddEntitiesCallback

from . import AutomowerConfigEntry
from .coordinator import AutomowerDataUpdateCoordinator
from .entity import AutomowerControlEntity, handle_sending_exception
from .entity import (
AutomowerControlEntity,
WorkAreaControlEntity,
_work_area_translation_key,
async_remove_work_area_entities,
handle_sending_exception,
)

_LOGGER = logging.getLogger(__name__)

Expand All @@ -30,14 +35,6 @@ def _async_get_cutting_height(data: MowerAttributes) -> int:
return data.settings.cutting_height


@callback
def _work_area_translation_key(work_area_id: int) -> str:
"""Return the translation key."""
if work_area_id == 0:
return "my_lawn_cutting_height"
return "work_area_cutting_height"


async def async_set_work_area_cutting_height(
coordinator: AutomowerDataUpdateCoordinator,
mower_id: str,
Expand Down Expand Up @@ -88,7 +85,7 @@ class AutomowerWorkAreaNumberEntityDescription(NumberEntityDescription):
"""Describes Automower work area number entity."""

value_fn: Callable[[WorkArea], int]
translation_key_fn: Callable[[int], str]
translation_key_fn: Callable[[int, str], str]
set_value_fn: Callable[
[AutomowerDataUpdateCoordinator, str, float, int], Awaitable[Any]
]
Expand Down Expand Up @@ -126,7 +123,7 @@ async def async_setup_entry(
for description in WORK_AREA_NUMBER_TYPES
for work_area_id in _work_areas
)
async_remove_entities(hass, coordinator, entry, mower_id)
async_remove_work_area_entities(hass, coordinator, entry, mower_id)
entities.extend(
AutomowerNumberEntity(mower_id, coordinator, description)
for description in NUMBER_TYPES
Expand Down Expand Up @@ -164,7 +161,7 @@ async def async_set_native_value(self, value: float) -> None:
)


class AutomowerWorkAreaNumberEntity(AutomowerControlEntity, NumberEntity):
class AutomowerWorkAreaNumberEntity(WorkAreaControlEntity, NumberEntity):
"""Defining the AutomowerWorkAreaNumberEntity with AutomowerWorkAreaNumberEntityDescription."""

entity_description: AutomowerWorkAreaNumberEntityDescription
Expand All @@ -177,57 +174,28 @@ def __init__(
work_area_id: int,
) -> None:
"""Set up AutomowerNumberEntity."""
super().__init__(mower_id, coordinator)
super().__init__(mower_id, coordinator, work_area_id)
self.entity_description = description
self.work_area_id = work_area_id
self._attr_unique_id = f"{mower_id}_{work_area_id}_{description.key}"
self._attr_translation_placeholders = {"work_area": self.work_area.name}

@property
def work_area(self) -> WorkArea:
"""Get the mower attributes of the current mower."""
if TYPE_CHECKING:
assert self.mower_attributes.work_areas is not None
return self.mower_attributes.work_areas[self.work_area_id]
self._attr_translation_placeholders = {
"work_area": self.work_area_attributes.name
}

@property
def translation_key(self) -> str:
"""Return the translation key of the work area."""
return self.entity_description.translation_key_fn(self.work_area_id)
return self.entity_description.translation_key_fn(
self.work_area_id, self.entity_description.key
)

@property
def native_value(self) -> float:
"""Return the state of the number."""
return self.entity_description.value_fn(self.work_area)
return self.entity_description.value_fn(self.work_area_attributes)

@handle_sending_exception(poll_after_sending=True)
async def async_set_native_value(self, value: float) -> None:
"""Change to new number value."""
await self.entity_description.set_value_fn(
self.coordinator, self.mower_id, value, self.work_area_id
)


@callback
def async_remove_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 (
entity_entry.domain == Platform.NUMBER
and (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)
7 changes: 5 additions & 2 deletions homeassistant/components/husqvarna_automower/strings.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,10 +54,10 @@
"cutting_height": {
"name": "Cutting height"
},
"my_lawn_cutting_height": {
"my_lawn_cutting_height_work_area": {
"name": "My lawn cutting height"
},
"work_area_cutting_height": {
"work_area_cutting_height_work_area": {
"name": "{work_area} cutting height"
}
},
Expand Down Expand Up @@ -271,6 +271,9 @@
},
"stay_out_zones": {
"name": "Avoid {stay_out_zone}"
},
"my_lawn_work_area": {
"name": "My lawn"
}
}
},
Expand Down
55 changes: 54 additions & 1 deletion homeassistant/components/husqvarna_automower/switch.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,12 @@

from . import AutomowerConfigEntry
from .coordinator import AutomowerDataUpdateCoordinator
from .entity import AutomowerControlEntity, handle_sending_exception
from .entity import (
AutomowerControlEntity,
WorkAreaControlEntity,
_work_area_translation_key,
handle_sending_exception,
)

_LOGGER = logging.getLogger(__name__)

Expand Down Expand Up @@ -41,6 +46,13 @@ async def async_setup_entry(
for stay_out_zone_uid in _stay_out_zones.zones
)
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
)
async_add_entities(entities)


Expand Down Expand Up @@ -131,6 +143,47 @@ async def async_turn_on(self, **kwargs: Any) -> None:
)


class WorkAreaSwitchEntity(WorkAreaControlEntity, SwitchEntity):
"""Defining the Automower work area switch."""

def __init__(
self,
coordinator: AutomowerDataUpdateCoordinator,
mower_id: str,
work_area_id: int,
) -> None:
"""Set up Automower switch."""
super().__init__(mower_id, coordinator, work_area_id)
key = "work_area"
self._attr_translation_key = _work_area_translation_key(work_area_id, key)
self._attr_unique_id = f"{mower_id}_{work_area_id}_{key}"
if self.work_area_attributes.name == "my_lawn":
self._attr_translation_placeholders = {
"work_area": self.work_area_attributes.name
}
else:
self._attr_name = self.work_area_attributes.name

@property
def is_on(self) -> bool:
"""Return the state of the switch."""
return self.work_area_attributes.enabled

@handle_sending_exception(poll_after_sending=True)
async def async_turn_off(self, **kwargs: Any) -> None:
"""Turn the switch off."""
await self.coordinator.api.commands.workarea_settings(
self.mower_id, self.work_area_id, enabled=False
)

@handle_sending_exception(poll_after_sending=True)
async def async_turn_on(self, **kwargs: Any) -> None:
"""Turn the switch on."""
await self.coordinator.api.commands.workarea_settings(
self.mower_id, self.work_area_id, enabled=True
)


@callback
def async_remove_entities(
hass: HomeAssistant,
Expand Down
6 changes: 3 additions & 3 deletions tests/components/husqvarna_automower/snapshots/test_number.ambr
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
'platform': 'husqvarna_automower',
'previous_unique_id': None,
'supported_features': 0,
'translation_key': 'work_area_cutting_height',
'translation_key': 'work_area_cutting_height_work_area',
'unique_id': 'c7233734-b219-4287-a173-08e3643f89f0_654321_cutting_height_work_area',
'unit_of_measurement': '%',
})
Expand Down Expand Up @@ -143,7 +143,7 @@
'platform': 'husqvarna_automower',
'previous_unique_id': None,
'supported_features': 0,
'translation_key': 'work_area_cutting_height',
'translation_key': 'work_area_cutting_height_work_area',
'unique_id': 'c7233734-b219-4287-a173-08e3643f89f0_123456_cutting_height_work_area',
'unit_of_measurement': '%',
})
Expand Down Expand Up @@ -199,7 +199,7 @@
'platform': 'husqvarna_automower',
'previous_unique_id': None,
'supported_features': 0,
'translation_key': 'my_lawn_cutting_height',
'translation_key': 'my_lawn_cutting_height_work_area',
'unique_id': 'c7233734-b219-4287-a173-08e3643f89f0_0_cutting_height_work_area',
'unit_of_measurement': '%',
})
Expand Down
Loading

0 comments on commit dc77b2d

Please sign in to comment.