Add Workarea cutting height to Husqvarna Automower (#116115)

* add work_area cutting_height

* add

* add default work_area

* ruff/mypy

* better names

* fit to api bump

* tweaks

* more tweaks

* layout

* address review

* change entity name

* tweak test

* cleanup entities

* fix for mowers with no workareas

* assure not other entities get deleted

* sort & remove one callback

* remove typing callbacks

* rename entity to entity_entry

homeassistant/components/husqvarna_automower/number.py

@@ -1,19 +1,21 @@
 """Creates the number entities for the mower."""
 
+import asyncio
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 import logging
 from typing import TYPE_CHECKING, Any
 
 from aioautomower.exceptions import ApiException
-from aioautomower.model import MowerAttributes
+from aioautomower.model import MowerAttributes, WorkArea
 from aioautomower.session import AutomowerSession
 
 from homeassistant.components.number import NumberEntity, NumberEntityDescription
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import EntityCategory
+from homeassistant.const import PERCENTAGE, EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
@@ -23,15 +25,6 @@ from .entity import AutomowerBaseEntity
 _LOGGER = logging.getLogger(__name__)
 
 
-@dataclass(frozen=True, kw_only=True)
-class AutomowerNumberEntityDescription(NumberEntityDescription):
-    """Describes Automower number entity."""
-
-    exists_fn: Callable[[MowerAttributes], bool] = lambda _: True
-    value_fn: Callable[[MowerAttributes], int]
-    set_value_fn: Callable[[AutomowerSession, str, float], Awaitable[Any]]
-
-
 @callback
 def _async_get_cutting_height(data: MowerAttributes) -> int:
     """Return the cutting height."""
@@ -41,6 +34,39 @@ def _async_get_cutting_height(data: MowerAttributes) -> int:
     return data.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,
+    cheight: float,
+    work_area_id: int,
+) -> None:
+    """Set cutting height for work area."""
+    await coordinator.api.set_cutting_height_workarea(
+        mower_id, int(cheight), work_area_id
+    )
+    # As there are no updates from the websocket regarding work area changes,
+    # we need to wait 5s and then poll the API.
+    await asyncio.sleep(5)
+    await coordinator.async_request_refresh()
+
+
+@dataclass(frozen=True, kw_only=True)
+class AutomowerNumberEntityDescription(NumberEntityDescription):
+    """Describes Automower number entity."""
+
+    exists_fn: Callable[[MowerAttributes], bool] = lambda _: True
+    value_fn: Callable[[MowerAttributes], int]
+    set_value_fn: Callable[[AutomowerSession, str, float], Awaitable[Any]]
+
+
 NUMBER_TYPES: tuple[AutomowerNumberEntityDescription, ...] = (
     AutomowerNumberEntityDescription(
         key="cutting_height",
@@ -58,17 +84,55 @@ NUMBER_TYPES: tuple[AutomowerNumberEntityDescription, ...] = (
 )
 
 
+@dataclass(frozen=True, kw_only=True)
+class AutomowerWorkAreaNumberEntityDescription(NumberEntityDescription):
+    """Describes Automower work area number entity."""
+
+    value_fn: Callable[[WorkArea], int]
+    translation_key_fn: Callable[[int], str]
+    set_value_fn: Callable[
+        [AutomowerDataUpdateCoordinator, str, float, int], Awaitable[Any]
+    ]
+
+
+WORK_AREA_NUMBER_TYPES: tuple[AutomowerWorkAreaNumberEntityDescription, ...] = (
+    AutomowerWorkAreaNumberEntityDescription(
+        key="cutting_height_work_area",
+        translation_key_fn=_work_area_translation_key,
+        entity_category=EntityCategory.CONFIG,
+        native_unit_of_measurement=PERCENTAGE,
+        value_fn=lambda data: data.cutting_height,
+        set_value_fn=async_set_work_area_cutting_height,
+    ),
+)
+
+
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up number platform."""
     coordinator: AutomowerDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities(
+    entities: list[NumberEntity] = []
+
+    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(
+                    AutomowerWorkAreaNumberEntity(
+                        mower_id, coordinator, description, work_area_id
+                    )
+                    for description in WORK_AREA_NUMBER_TYPES
+                    for work_area_id in _work_areas
+                )
+            await async_remove_entities(coordinator, hass, entry, mower_id)
+    entities.extend(
         AutomowerNumberEntity(mower_id, coordinator, description)
         for mower_id in coordinator.data
         for description in NUMBER_TYPES
         if description.exists_fn(coordinator.data[mower_id])
     )
+    async_add_entities(entities)
 
 
 class AutomowerNumberEntity(AutomowerBaseEntity, NumberEntity):
@@ -102,3 +166,74 @@ class AutomowerNumberEntity(AutomowerBaseEntity, NumberEntity):
             raise HomeAssistantError(
                 f"Command couldn't be sent to the command queue: {exception}"
             ) from exception
+
+
+class AutomowerWorkAreaNumberEntity(AutomowerBaseEntity, NumberEntity):
+    """Defining the AutomowerWorkAreaNumberEntity with AutomowerWorkAreaNumberEntityDescription."""
+
+    entity_description: AutomowerWorkAreaNumberEntityDescription
+
+    def __init__(
+        self,
+        mower_id: str,
+        coordinator: AutomowerDataUpdateCoordinator,
+        description: AutomowerWorkAreaNumberEntityDescription,
+        work_area_id: int,
+    ) -> None:
+        """Set up AutomowerNumberEntity."""
+        super().__init__(mower_id, coordinator)
+        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]
+
+    @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)
+
+    @property
+    def native_value(self) -> float:
+        """Return the state of the number."""
+        return self.entity_description.value_fn(self.work_area)
+
+    async def async_set_native_value(self, value: float) -> None:
+        """Change to new number value."""
+        try:
+            await self.entity_description.set_value_fn(
+                self.coordinator, self.mower_id, value, self.work_area_id
+            )
+        except ApiException as exception:
+            raise HomeAssistantError(
+                f"Command couldn't be sent to the command queue: {exception}"
+            ) from exception
+
+
+async def async_remove_entities(
+    coordinator: AutomowerDataUpdateCoordinator,
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    mower_id: str,
+) -> None:
+    """Remove deleted work areas from Home Assistant."""
+    entity_reg = er.async_get(hass)
+    work_area_list = []
+    _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"
+            work_area_list.append(uid)
+        for entity_entry in er.async_entries_for_config_entry(
+            entity_reg, config_entry.entry_id
+        ):
+            if entity_entry.unique_id.split("_")[0] == mower_id:
+                if entity_entry.unique_id.endswith("cutting_height_work_area"):
+                    if entity_entry.unique_id not in work_area_list:
+                        entity_reg.async_remove(entity_entry.entity_id)

homeassistant/components/husqvarna_automower/strings.json

@@ -40,6 +40,12 @@
     "number": {
       "cutting_height": {
         "name": "Cutting height"
+      },
+      "my_lawn_cutting_height": {
+        "name": "My lawn cutting height "
+      },
+      "work_area_cutting_height": {
+        "name": "{work_area} cutting height"
       }
     },
     "select": {

tests/components/husqvarna_automower/fixtures/mower.json

@@ -14,9 +14,9 @@
         },
         "capabilities": {
           "headlights": true,
-          "workAreas": false,
+          "workAreas": true,
           "position": true,
-          "stayOutZones": false
+          "stayOutZones": true
         },
         "mower": {
           "mode": "MAIN_AREA",
@@ -68,6 +68,11 @@
             "name": "Front lawn",
             "cuttingHeight": 50
           },
+          {
+            "workAreaId": 654321,
+            "name": "Back lawn",
+            "cuttingHeight": 25
+          },
           {
             "workAreaId": 0,
             "name": "",

tests/components/husqvarna_automower/snapshots/test_diagnostics.ambr

@@ -35,8 +35,8 @@
     'capabilities': dict({
       'headlights': True,
       'position': True,
-      'stay_out_zones': False,
+      'stay_out_zones': True,
-      'work_areas': False,
+      'work_areas': True,
     }),
     'cutting_height': 4,
     'headlight': dict({
@@ -97,6 +97,10 @@
         'cutting_height': 50,
         'name': 'Front lawn',
       }),
+      '654321': dict({
+        'cutting_height': 25,
+        'name': 'Back lawn',
+      }),
     }),
   })
 # ---

tests/components/husqvarna_automower/snapshots/test_number.ambr

@@ -1,4 +1,60 @@
 # serializer version: 1
+# name: test_snapshot_number[number.test_mower_1_back_lawn_cutting_height-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': dict({
+      'max': 100.0,
+      'min': 0.0,
+      'mode': <NumberMode.AUTO: 'auto'>,
+      'step': 1.0,
+    }),
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'number',
+    'entity_category': <EntityCategory.CONFIG: 'config'>,
+    'entity_id': 'number.test_mower_1_back_lawn_cutting_height',
+    'has_entity_name': True,
+    'hidden_by': None,
+    'icon': None,
+    'id': <ANY>,
+    'labels': set({
+    }),
+    'name': None,
+    'options': dict({
+    }),
+    'original_device_class': None,
+    'original_icon': None,
+    'original_name': 'Back lawn cutting height',
+    'platform': 'husqvarna_automower',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': 'work_area_cutting_height',
+    'unique_id': 'c7233734-b219-4287-a173-08e3643f89f0_654321_cutting_height_work_area',
+    'unit_of_measurement': '%',
+  })
+# ---
+# name: test_snapshot_number[number.test_mower_1_back_lawn_cutting_height-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Mower 1 Back lawn cutting height',
+      'max': 100.0,
+      'min': 0.0,
+      'mode': <NumberMode.AUTO: 'auto'>,
+      'step': 1.0,
+      'unit_of_measurement': '%',
+    }),
+    'context': <ANY>,
+    'entity_id': 'number.test_mower_1_back_lawn_cutting_height',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': '25',
+  })
+# ---
 # name: test_snapshot_number[number.test_mower_1_cutting_height-entry]
   EntityRegistryEntrySnapshot({
     'aliases': set({
@@ -54,3 +110,115 @@
     'state': '4',
   })
 # ---
+# name: test_snapshot_number[number.test_mower_1_front_lawn_cutting_height-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': dict({
+      'max': 100.0,
+      'min': 0.0,
+      'mode': <NumberMode.AUTO: 'auto'>,
+      'step': 1.0,
+    }),
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'number',
+    'entity_category': <EntityCategory.CONFIG: 'config'>,
+    'entity_id': 'number.test_mower_1_front_lawn_cutting_height',
+    'has_entity_name': True,
+    'hidden_by': None,
+    'icon': None,
+    'id': <ANY>,
+    'labels': set({
+    }),
+    'name': None,
+    'options': dict({
+    }),
+    'original_device_class': None,
+    'original_icon': None,
+    'original_name': 'Front lawn cutting height',
+    'platform': 'husqvarna_automower',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': 'work_area_cutting_height',
+    'unique_id': 'c7233734-b219-4287-a173-08e3643f89f0_123456_cutting_height_work_area',
+    'unit_of_measurement': '%',
+  })
+# ---
+# name: test_snapshot_number[number.test_mower_1_front_lawn_cutting_height-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Mower 1 Front lawn cutting height',
+      'max': 100.0,
+      'min': 0.0,
+      'mode': <NumberMode.AUTO: 'auto'>,
+      'step': 1.0,
+      'unit_of_measurement': '%',
+    }),
+    'context': <ANY>,
+    'entity_id': 'number.test_mower_1_front_lawn_cutting_height',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': '50',
+  })
+# ---
+# name: test_snapshot_number[number.test_mower_1_my_lawn_cutting_height-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': dict({
+      'max': 100.0,
+      'min': 0.0,
+      'mode': <NumberMode.AUTO: 'auto'>,
+      'step': 1.0,
+    }),
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'number',
+    'entity_category': <EntityCategory.CONFIG: 'config'>,
+    'entity_id': 'number.test_mower_1_my_lawn_cutting_height',
+    'has_entity_name': True,
+    'hidden_by': None,
+    'icon': None,
+    'id': <ANY>,
+    'labels': set({
+    }),
+    'name': None,
+    'options': dict({
+    }),
+    'original_device_class': None,
+    'original_icon': None,
+    'original_name': 'My lawn cutting height ',
+    'platform': 'husqvarna_automower',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': 'my_lawn_cutting_height',
+    'unique_id': 'c7233734-b219-4287-a173-08e3643f89f0_0_cutting_height_work_area',
+    'unit_of_measurement': '%',
+  })
+# ---
+# name: test_snapshot_number[number.test_mower_1_my_lawn_cutting_height-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Mower 1 My lawn cutting height ',
+      'max': 100.0,
+      'min': 0.0,
+      'mode': <NumberMode.AUTO: 'auto'>,
+      'step': 1.0,
+      'unit_of_measurement': '%',
+    }),
+    'context': <ANY>,
+    'entity_id': 'number.test_mower_1_my_lawn_cutting_height',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': '50',
+  })
+# ---

tests/components/husqvarna_automower/test_number.py

@@ -3,17 +3,20 @@
 from unittest.mock import AsyncMock, patch
 
 from aioautomower.exceptions import ApiException
+from aioautomower.utils import mower_list_to_dictionary_dataclass
 import pytest
 from syrupy import SnapshotAssertion
 
+from homeassistant.components.husqvarna_automower.const import DOMAIN
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
 
 from . import setup_integration
+from .const import TEST_MOWER_ID
 
-from tests.common import MockConfigEntry, snapshot_platform
+from tests.common import MockConfigEntry, load_json_value_fixture, snapshot_platform
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
@@ -51,6 +54,74 @@ async def test_number_commands(
     assert len(mocked_method.mock_calls) == 2
 
 
+async def test_number_workarea_commands(
+    hass: HomeAssistant,
+    mock_automower_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test number commands."""
+    entity_id = "number.test_mower_1_front_lawn_cutting_height"
+    await setup_integration(hass, mock_config_entry)
+    values = mower_list_to_dictionary_dataclass(
+        load_json_value_fixture("mower.json", DOMAIN)
+    )
+    values[TEST_MOWER_ID].work_areas[123456].cutting_height = 75
+    mock_automower_client.get_status.return_value = values
+    mocked_method = AsyncMock()
+    setattr(mock_automower_client, "set_cutting_height_workarea", mocked_method)
+    await hass.services.async_call(
+        domain="number",
+        service="set_value",
+        target={"entity_id": entity_id},
+        service_data={"value": "75"},
+        blocking=True,
+    )
+    assert len(mocked_method.mock_calls) == 1
+    state = hass.states.get(entity_id)
+    assert state.state is not None
+    assert state.state == "75"
+
+    mocked_method.side_effect = ApiException("Test error")
+    with pytest.raises(HomeAssistantError) as exc_info:
+        await hass.services.async_call(
+            domain="number",
+            service="set_value",
+            target={"entity_id": entity_id},
+            service_data={"value": "75"},
+            blocking=True,
+        )
+    assert (
+        str(exc_info.value)
+        == "Command couldn't be sent to the command queue: Test error"
+    )
+    assert len(mocked_method.mock_calls) == 2
+
+
+async def test_workarea_deleted(
+    hass: HomeAssistant,
+    mock_automower_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    entity_registry: er.EntityRegistry,
+) -> None:
+    """Test if work area is deleted after removed."""
+
+    values = mower_list_to_dictionary_dataclass(
+        load_json_value_fixture("mower.json", DOMAIN)
+    )
+    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 - 1)
+
+
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
 async def test_snapshot_number(
     hass: HomeAssistant,