Add work area switch for Husqvarna Automower (#126376)

* 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 <joostlek@outlook.com>

* constants

---------

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

homeassistant/components/husqvarna_automower/entity.py

@@ -4,17 +4,18 @@ import asyncio
 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__)
@@ -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[
@@ -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

homeassistant/components/husqvarna_automower/number.py

@@ -9,14 +9,19 @@ from aioautomower.model import MowerAttributes, WorkArea
 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__)
 
@@ -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,
@@ -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]
     ]
@@ -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
@@ -164,7 +161,7 @@ class AutomowerNumberEntity(AutomowerControlEntity, NumberEntity):
         )
 
 
-class AutomowerWorkAreaNumberEntity(AutomowerControlEntity, NumberEntity):
+class AutomowerWorkAreaNumberEntity(WorkAreaControlEntity, NumberEntity):
     """Defining the AutomowerWorkAreaNumberEntity with AutomowerWorkAreaNumberEntityDescription."""
 
     entity_description: AutomowerWorkAreaNumberEntityDescription
@@ -177,28 +174,24 @@ class AutomowerWorkAreaNumberEntity(AutomowerControlEntity, NumberEntity):
         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}
+        self._attr_translation_placeholders = {
-
+            "work_area": self.work_area_attributes.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)
+        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:
@@ -206,28 +199,3 @@ class AutomowerWorkAreaNumberEntity(AutomowerControlEntity, NumberEntity):
         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)

homeassistant/components/husqvarna_automower/strings.json

@@ -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"
       }
     },
@@ -271,6 +271,9 @@
       },
       "stay_out_zones": {
         "name": "Avoid {stay_out_zone}"
+      },
+      "my_lawn_work_area": {
+        "name": "My lawn"
       }
     }
   },

homeassistant/components/husqvarna_automower/switch.py

@@ -13,7 +13,12 @@ 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,
+    handle_sending_exception,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -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)
 
 
@@ -131,6 +143,47 @@ class AutomowerStayOutZoneSwitchEntity(AutomowerControlEntity, SwitchEntity):
         )
 
 
+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,

tests/components/husqvarna_automower/snapshots/test_number.ambr

@@ -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': '%',
   })
@@ -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': '%',
   })
@@ -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': '%',
   })

tests/components/husqvarna_automower/snapshots/test_switch.ambr

@@ -91,6 +91,52 @@
     'state': 'on',
   })
 # ---
+# name: test_switch_snapshot[switch.test_mower_1_back_lawn-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': None,
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'switch',
+    'entity_category': None,
+    'entity_id': 'switch.test_mower_1_back_lawn',
+    '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',
+    'platform': 'husqvarna_automower',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': 'work_area_work_area',
+    'unique_id': 'c7233734-b219-4287-a173-08e3643f89f0_654321_work_area',
+    'unit_of_measurement': None,
+  })
+# ---
+# name: test_switch_snapshot[switch.test_mower_1_back_lawn-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Mower 1 Back lawn',
+    }),
+    'context': <ANY>,
+    'entity_id': 'switch.test_mower_1_back_lawn',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on',
+  })
+# ---
 # name: test_switch_snapshot[switch.test_mower_1_enable_schedule-entry]
   EntityRegistryEntrySnapshot({
     'aliases': set({
@@ -137,6 +183,98 @@
     'state': 'on',
   })
 # ---
+# name: test_switch_snapshot[switch.test_mower_1_front_lawn-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': None,
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'switch',
+    'entity_category': None,
+    'entity_id': 'switch.test_mower_1_front_lawn',
+    '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',
+    'platform': 'husqvarna_automower',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': 'work_area_work_area',
+    'unique_id': 'c7233734-b219-4287-a173-08e3643f89f0_123456_work_area',
+    'unit_of_measurement': None,
+  })
+# ---
+# name: test_switch_snapshot[switch.test_mower_1_front_lawn-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Mower 1 Front lawn',
+    }),
+    'context': <ANY>,
+    'entity_id': 'switch.test_mower_1_front_lawn',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on',
+  })
+# ---
+# name: test_switch_snapshot[switch.test_mower_1_my_lawn-entry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': None,
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'switch',
+    'entity_category': None,
+    'entity_id': 'switch.test_mower_1_my_lawn',
+    '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',
+    'platform': 'husqvarna_automower',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': 'my_lawn_work_area',
+    'unique_id': 'c7233734-b219-4287-a173-08e3643f89f0_0_work_area',
+    'unit_of_measurement': None,
+  })
+# ---
+# name: test_switch_snapshot[switch.test_mower_1_my_lawn-state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Test Mower 1 My lawn',
+    }),
+    'context': <ANY>,
+    'entity_id': 'switch.test_mower_1_my_lawn',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'off',
+  })
+# ---
 # name: test_switch_snapshot[switch.test_mower_2_enable_schedule-entry]
   EntityRegistryEntrySnapshot({
     'aliases': set({

tests/components/husqvarna_automower/test_init.py

@@ -167,6 +167,31 @@ 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,
+) -> 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 - 2)
+
+
 async def test_coordinator_automatic_registry_cleanup(
     hass: HomeAssistant,
     mock_automower_client: AsyncMock,
@@ -179,8 +204,12 @@ async def test_coordinator_automatic_registry_cleanup(
     entry = hass.config_entries.async_entries(DOMAIN)[0]
     await hass.async_block_till_done()
 
-    assert len(er.async_entries_for_config_entry(entity_registry, entry.entry_id)) == 42
+    current_entites = len(
-    assert len(dr.async_entries_for_config_entry(device_registry, entry.entry_id)) == 2
+        er.async_entries_for_config_entry(entity_registry, entry.entry_id)
+    )
+    current_devices = len(
+        dr.async_entries_for_config_entry(device_registry, entry.entry_id)
+    )
 
     values = mower_list_to_dictionary_dataclass(
         load_json_value_fixture("mower.json", DOMAIN)
@@ -190,5 +219,11 @@ async def test_coordinator_automatic_registry_cleanup(
     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, entry.entry_id)) == 12
+    assert (
-    assert len(dr.async_entries_for_config_entry(device_registry, entry.entry_id)) == 1
+        len(er.async_entries_for_config_entry(entity_registry, entry.entry_id))
+        == current_entites - 33
+    )
+    assert (
+        len(dr.async_entries_for_config_entry(device_registry, entry.entry_id))
+        == current_devices - 1
+    )

tests/components/husqvarna_automower/test_number.py

@@ -109,31 +109,6 @@ async def test_number_workarea_commands(
     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_number_snapshot(
     hass: HomeAssistant,

tests/components/husqvarna_automower/test_switch.py

@@ -15,7 +15,14 @@ from homeassistant.components.husqvarna_automower.const import (
     EXECUTION_TIME_DELAY,
 )
 from homeassistant.components.husqvarna_automower.coordinator import SCAN_INTERVAL
-from homeassistant.const import Platform
+from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    SERVICE_TOGGLE,
+    SERVICE_TURN_OFF,
+    SERVICE_TURN_ON,
+    Platform,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
@@ -31,6 +38,7 @@ from tests.common import (
 )
 
 TEST_ZONE_ID = "AAAAAAAA-BBBB-CCCC-DDDD-123456789101"
+TEST_AREA_ID = 0
 
 
 async def test_switch_states(
@@ -61,9 +69,9 @@ async def test_switch_states(
 @pytest.mark.parametrize(
     ("service", "aioautomower_command"),
     [
-        ("turn_off", "park_until_further_notice"),
+        (SERVICE_TURN_OFF, "park_until_further_notice"),
-        ("turn_on", "resume_schedule"),
+        (SERVICE_TURN_ON, "resume_schedule"),
-        ("toggle", "park_until_further_notice"),
+        (SERVICE_TOGGLE, "park_until_further_notice"),
     ],
 )
 async def test_switch_commands(
@@ -76,9 +84,9 @@ async def test_switch_commands(
     """Test switch commands."""
     await setup_integration(hass, mock_config_entry)
     await hass.services.async_call(
-        domain="switch",
+        domain=SWITCH_DOMAIN,
         service=service,
-        service_data={"entity_id": "switch.test_mower_1_enable_schedule"},
+        service_data={ATTR_ENTITY_ID: "switch.test_mower_1_enable_schedule"},
         blocking=True,
     )
     mocked_method = getattr(mock_automower_client.commands, aioautomower_command)
@@ -90,9 +98,9 @@ async def test_switch_commands(
         match="Failed to send command: Test error",
     ):
         await hass.services.async_call(
-            domain="switch",
+            domain=SWITCH_DOMAIN,
             service=service,
-            service_data={"entity_id": "switch.test_mower_1_enable_schedule"},
+            service_data={ATTR_ENTITY_ID: "switch.test_mower_1_enable_schedule"},
             blocking=True,
         )
     assert len(mocked_method.mock_calls) == 2
@@ -101,9 +109,9 @@ async def test_switch_commands(
 @pytest.mark.parametrize(
     ("service", "boolean", "excepted_state"),
     [
-        ("turn_off", False, "off"),
+        (SERVICE_TURN_OFF, False, "off"),
-        ("turn_on", True, "on"),
+        (SERVICE_TURN_ON, True, "on"),
-        ("toggle", True, "on"),
+        (SERVICE_TOGGLE, True, "on"),
     ],
 )
 async def test_stay_out_zone_switch_commands(
@@ -126,9 +134,9 @@ async def test_stay_out_zone_switch_commands(
     mocked_method = AsyncMock()
     setattr(mock_automower_client.commands, "switch_stay_out_zone", mocked_method)
     await hass.services.async_call(
-        domain="switch",
+        domain=SWITCH_DOMAIN,
         service=service,
-        service_data={"entity_id": entity_id},
+        service_data={ATTR_ENTITY_ID: entity_id},
         blocking=False,
     )
     freezer.tick(timedelta(seconds=EXECUTION_TIME_DELAY))
@@ -145,9 +153,64 @@ async def test_stay_out_zone_switch_commands(
         match="Failed to send command: Test error",
     ):
         await hass.services.async_call(
-            domain="switch",
+            domain=SWITCH_DOMAIN,
             service=service,
-            service_data={"entity_id": entity_id},
+            service_data={ATTR_ENTITY_ID: entity_id},
+            blocking=True,
+        )
+    assert len(mocked_method.mock_calls) == 2
+
+
+@pytest.mark.parametrize(
+    ("service", "boolean", "excepted_state"),
+    [
+        (SERVICE_TURN_OFF, False, "off"),
+        (SERVICE_TURN_ON, True, "on"),
+        (SERVICE_TOGGLE, True, "on"),
+    ],
+)
+async def test_work_area_switch_commands(
+    hass: HomeAssistant,
+    service: str,
+    boolean: bool,
+    excepted_state: str,
+    mock_automower_client: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test switch commands."""
+    entity_id = "switch.test_mower_1_my_lawn"
+    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[TEST_AREA_ID].enabled = boolean
+    mock_automower_client.get_status.return_value = values
+    mocked_method = AsyncMock()
+    setattr(mock_automower_client.commands, "workarea_settings", mocked_method)
+    await hass.services.async_call(
+        domain=SWITCH_DOMAIN,
+        service=service,
+        service_data={ATTR_ENTITY_ID: entity_id},
+        blocking=False,
+    )
+    freezer.tick(timedelta(seconds=EXECUTION_TIME_DELAY))
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+    mocked_method.assert_called_once_with(TEST_MOWER_ID, TEST_AREA_ID, enabled=boolean)
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == excepted_state
+
+    mocked_method.side_effect = ApiException("Test error")
+    with pytest.raises(
+        HomeAssistantError,
+        match="Failed to send command: Test error",
+    ):
+        await hass.services.async_call(
+            domain=SWITCH_DOMAIN,
+            service=service,
+            service_data={ATTR_ENTITY_ID: entity_id},
             blocking=True,
         )
     assert len(mocked_method.mock_calls) == 2