Add override for work areas in Husqvarna Automower (#120427)

Co-authored-by: Robert Resch <robert@resch.dev>

homeassistant/components/husqvarna_automower/icons.json

@@ -34,6 +34,7 @@
     }
   },
   "services": {
-    "override_schedule": "mdi:debug-step-over"
+    "override_schedule": "mdi:debug-step-over",
+    "override_schedule_work_area": "mdi:land-fields"
   }
 }

homeassistant/components/husqvarna_automower/lawn_mower.py

@@ -2,8 +2,9 @@
 
 from datetime import timedelta
 import logging
+from typing import TYPE_CHECKING
 
-from aioautomower.model import MowerActivities, MowerStates
+from aioautomower.model import MowerActivities, MowerStates, WorkArea
 import voluptuous as vol
 
 from homeassistant.components.lawn_mower import (
@@ -12,10 +13,12 @@ from homeassistant.components.lawn_mower import (
     LawnMowerEntityFeature,
 )
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ServiceValidationError
 from homeassistant.helpers import config_validation as cv, entity_platform
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import AutomowerConfigEntry
+from .const import DOMAIN
 from .coordinator import AutomowerDataUpdateCoordinator
 from .entity import AutomowerAvailableEntity, handle_sending_exception
 
@@ -67,6 +70,18 @@ async def async_setup_entry(
         },
         "async_override_schedule",
     )
+    platform.async_register_entity_service(
+        "override_schedule_work_area",
+        {
+            vol.Required("work_area_id"): vol.Coerce(int),
+            vol.Required("duration"): vol.All(
+                cv.time_period,
+                cv.positive_timedelta,
+                vol.Range(min=timedelta(minutes=1), max=timedelta(days=42)),
+            ),
+        },
+        "async_override_schedule_work_area",
+    )
 
 
 class AutomowerLawnMowerEntity(AutomowerAvailableEntity, LawnMowerEntity):
@@ -98,6 +113,11 @@ class AutomowerLawnMowerEntity(AutomowerAvailableEntity, LawnMowerEntity):
             return LawnMowerActivity.DOCKED
         return LawnMowerActivity.ERROR
 
+    @property
+    def work_areas(self) -> dict[int, WorkArea] | None:
+        """Return the work areas of the mower."""
+        return self.mower_attributes.work_areas
+
     @handle_sending_exception()
     async def async_start_mowing(self) -> None:
         """Resume schedule."""
@@ -122,3 +142,22 @@ class AutomowerLawnMowerEntity(AutomowerAvailableEntity, LawnMowerEntity):
             await self.coordinator.api.commands.start_for(self.mower_id, duration)
         if override_mode == PARK:
             await self.coordinator.api.commands.park_for(self.mower_id, duration)
+
+    @handle_sending_exception()
+    async def async_override_schedule_work_area(
+        self, work_area_id: int, duration: timedelta
+    ) -> None:
+        """Override the schedule with a certain work area."""
+        if not self.mower_attributes.capabilities.work_areas:
+            raise ServiceValidationError(
+                translation_domain=DOMAIN, translation_key="work_areas_not_supported"
+            )
+        if TYPE_CHECKING:
+            assert self.work_areas is not None
+        if work_area_id not in self.work_areas:
+            raise ServiceValidationError(
+                translation_domain=DOMAIN, translation_key="work_area_not_existing"
+            )
+        await self.coordinator.api.commands.start_in_workarea(
+            self.mower_id, work_area_id, duration
+        )

homeassistant/components/husqvarna_automower/sensor.py

@@ -1,10 +1,10 @@
 """Creates the sensor entities for the mower."""
 
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass
 from datetime import datetime
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 from zoneinfo import ZoneInfo
 
 from aioautomower.model import MowerAttributes, MowerModes, RestrictedReasons
@@ -28,6 +28,8 @@ from .entity import AutomowerBaseEntity
 
 _LOGGER = logging.getLogger(__name__)
 
+ATTR_WORK_AREA_ID_ASSIGNMENT = "work_area_id_assignment"
+
 ERROR_KEY_LIST = [
     "no_error",
     "alarm_mower_in_motion",
@@ -214,11 +216,23 @@ def _get_current_work_area_name(data: MowerAttributes) -> str:
     return data.work_areas[data.mower.work_area_id].name
 
 
+@callback
+def _get_current_work_area_dict(data: MowerAttributes) -> Mapping[str, Any]:
+    """Return the name of the current work area."""
+    if TYPE_CHECKING:
+        # Sensor does not get created if it is None
+        assert data.work_areas is not None
+    return {ATTR_WORK_AREA_ID_ASSIGNMENT: data.work_area_dict}
+
+
 @dataclass(frozen=True, kw_only=True)
 class AutomowerSensorEntityDescription(SensorEntityDescription):
     """Describes Automower sensor entity."""
 
     exists_fn: Callable[[MowerAttributes], bool] = lambda _: True
+    extra_state_attributes_fn: Callable[[MowerAttributes], Mapping[str, Any] | None] = (
+        lambda _: None
+    )
     option_fn: Callable[[MowerAttributes], list[str] | None] = lambda _: None
     value_fn: Callable[[MowerAttributes], StateType | datetime]
 
@@ -353,6 +367,7 @@ SENSOR_TYPES: tuple[AutomowerSensorEntityDescription, ...] = (
         translation_key="work_area",
         device_class=SensorDeviceClass.ENUM,
         exists_fn=lambda data: data.capabilities.work_areas,
+        extra_state_attributes_fn=_get_current_work_area_dict,
         option_fn=_get_work_area_names,
         value_fn=_get_current_work_area_name,
     ),
@@ -378,6 +393,7 @@ class AutomowerSensorEntity(AutomowerBaseEntity, SensorEntity):
     """Defining the Automower Sensors with AutomowerSensorEntityDescription."""
 
     entity_description: AutomowerSensorEntityDescription
+    _unrecorded_attributes = frozenset({ATTR_WORK_AREA_ID_ASSIGNMENT})
 
     def __init__(
         self,
@@ -399,3 +415,8 @@ class AutomowerSensorEntity(AutomowerBaseEntity, SensorEntity):
     def options(self) -> list[str] | None:
         """Return the option of the sensor."""
         return self.entity_description.option_fn(self.mower_attributes)
+
+    @property
+    def extra_state_attributes(self) -> Mapping[str, Any] | None:
+        """Return the state attributes."""
+        return self.entity_description.extra_state_attributes_fn(self.mower_attributes)

homeassistant/components/husqvarna_automower/services.yaml

@@ -19,3 +19,22 @@ override_schedule:
           options:
             - "mow"
             - "park"
+
+override_schedule_work_area:
+  target:
+    entity:
+      integration: "husqvarna_automower"
+      domain: "lawn_mower"
+  fields:
+    duration:
+      required: true
+      example: "{'days': 1, 'hours': 12, 'minutes': 30}"
+      selector:
+        duration:
+          enable_day: true
+    work_area_id:
+      required: true
+      example: "123"
+      selector:
+        text:
+          type: number

homeassistant/components/husqvarna_automower/strings.json

@@ -254,6 +254,11 @@
         "state": {
           "my_lawn": "My lawn",
           "no_work_area_active": "No work area active"
+        },
+        "state_attributes": {
+          "work_area_id_assignment": {
+            "name": "Work area ID assignment"
+          }
         }
       }
     },
@@ -269,6 +274,12 @@
   "exceptions": {
     "command_send_failed": {
       "message": "Failed to send command: {exception}"
+    },
+    "work_areas_not_supported": {
+      "message": "This mower does not support work areas."
+    },
+    "work_area_not_existing": {
+      "message": "The selected work area does not exist."
     }
   },
   "selector": {
@@ -293,6 +304,20 @@
           "description": "With which action the schedule should be overridden."
         }
       }
+    },
+    "override_schedule_work_area": {
+      "name": "Override schedule work area",
+      "description": "Override the schedule of the mower for a duration of time in the selected work area.",
+      "fields": {
+        "duration": {
+          "name": "[%key:component::husqvarna_automower::services::override_schedule::fields::duration::name%]",
+          "description": "[%key:component::husqvarna_automower::services::override_schedule::fields::duration::description%]"
+        },
+        "work_area_id": {
+          "name": "Work area ID",
+          "description": "In which work area the mower should mow."
+        }
+      }
     }
   }
 }

tests/components/husqvarna_automower/snapshots/test_sensor.ambr

@@ -1100,6 +1100,11 @@
         'my_lawn',
         'no_work_area_active',
       ]),
+      'work_area_id_assignment': dict({
+        0: 'my_lawn',
+        123456: 'Front lawn',
+        654321: 'Back lawn',
+      }),
     }),
     'context': <ANY>,
     'entity_id': 'sensor.test_mower_1_work_area',

tests/components/husqvarna_automower/test_lawn_mower.py

@@ -13,7 +13,7 @@ from homeassistant.components.husqvarna_automower.const import DOMAIN
 from homeassistant.components.husqvarna_automower.coordinator import SCAN_INTERVAL
 from homeassistant.components.lawn_mower import LawnMowerActivity
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import HomeAssistantError
+from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
 
 from . import setup_integration
 from .const import TEST_MOWER_ID
@@ -122,7 +122,7 @@ async def test_lawn_mower_commands(
 async def test_lawn_mower_service_commands(
     hass: HomeAssistant,
     aioautomower_command: str,
-    extra_data: int | None,
+    extra_data: timedelta,
     service: str,
     service_data: dict[str, int] | None,
     mock_automower_client: AsyncMock,
@@ -158,27 +158,112 @@ async def test_lawn_mower_service_commands(
 
 
 @pytest.mark.parametrize(
-    ("service", "service_data"),
+    ("aioautomower_command", "extra_data1", "extra_data2", "service", "service_data"),
     [
         (
-            "override_schedule",
+            "start_in_workarea",
+            123456,
+            timedelta(days=40),
+            "override_schedule_work_area",
             {
-                "duration": {"days": 1, "hours": 12, "minutes": 30},
+                "work_area_id": 123456,
-                "override_mode": "fly_to_moon",
+                "duration": {"days": 40},
             },
         ),
     ],
 )
-async def test_lawn_mower_wrong_service_commands(
+async def test_lawn_mower_override_work_area_command(
     hass: HomeAssistant,
+    aioautomower_command: str,
+    extra_data1: int,
+    extra_data2: timedelta,
     service: str,
     service_data: dict[str, int] | None,
     mock_automower_client: AsyncMock,
     mock_config_entry: MockConfigEntry,
 ) -> None:
-    """Test lawn_mower commands."""
+    """Test lawn_mower work area override commands."""
     await setup_integration(hass, mock_config_entry)
-    with pytest.raises(MultipleInvalid):
+    mocked_method = AsyncMock()
+    setattr(mock_automower_client.commands, aioautomower_command, mocked_method)
+    await hass.services.async_call(
+        domain=DOMAIN,
+        service=service,
+        target={"entity_id": "lawn_mower.test_mower_1"},
+        service_data=service_data,
+        blocking=True,
+    )
+    mocked_method.assert_called_once_with(TEST_MOWER_ID, extra_data1, extra_data2)
+
+    getattr(
+        mock_automower_client.commands, aioautomower_command
+    ).side_effect = ApiException("Test error")
+    with pytest.raises(
+        HomeAssistantError,
+        match="Failed to send command: Test error",
+    ):
+        await hass.services.async_call(
+            domain=DOMAIN,
+            service=service,
+            target={"entity_id": "lawn_mower.test_mower_1"},
+            service_data=service_data,
+            blocking=True,
+        )
+
+
+@pytest.mark.parametrize(
+    ("service", "service_data", "mower_support_wa", "exception"),
+    [
+        (
+            "override_schedule",
+            {
+                "duration": {"days": 1, "hours": 12, "minutes": 30},
+                "override_mode": "fly_to_moon",
+            },
+            False,
+            MultipleInvalid,
+        ),
+        (
+            "override_schedule_work_area",
+            {
+                "work_area_id": 123456,
+                "duration": {"days": 40},
+            },
+            False,
+            ServiceValidationError,
+        ),
+        (
+            "override_schedule_work_area",
+            {
+                "work_area_id": 12345,
+                "duration": {"days": 40},
+            },
+            True,
+            ServiceValidationError,
+        ),
+    ],
+)
+async def test_lawn_mower_wrong_service_commands(
+    hass: HomeAssistant,
+    service: str,
+    service_data: dict[str, int] | None,
+    mower_support_wa: bool,
+    exception,
+    mock_automower_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test lawn_mower commands."""
+    await setup_integration(hass, mock_config_entry)
+    values = mower_list_to_dictionary_dataclass(
+        load_json_value_fixture("mower.json", DOMAIN)
+    )
+    values[TEST_MOWER_ID].capabilities.work_areas = mower_support_wa
+    mock_automower_client.get_status.return_value = values
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    with pytest.raises(exception):
         await hass.services.async_call(
             domain=DOMAIN,
             service=service,