Extend wrapper for sending commands to all platforms in Husqvarna Automower (#120255)

homeassistant/components/husqvarna_automower/entity.py

@@ -1,14 +1,20 @@
 """Platform for Husqvarna Automower base entity."""
 
+import asyncio
+from collections.abc import Awaitable, Callable, Coroutine
+import functools
 import logging
+from typing import Any
 
+from aioautomower.exceptions import ApiException
 from aioautomower.model import MowerActivities, MowerAttributes, MowerStates
 
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import AutomowerDataUpdateCoordinator
-from .const import DOMAIN
+from .const import DOMAIN, EXECUTION_TIME_DELAY
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -28,6 +34,38 @@ ERROR_STATES = [
 ]
 
 
+def handle_sending_exception(
+    poll_after_sending: bool = False,
+) -> Callable[
+    [Callable[..., Awaitable[Any]]], Callable[..., Coroutine[Any, Any, None]]
+]:
+    """Handle exceptions while sending a command and optionally refresh coordinator."""
+
+    def decorator(
+        func: Callable[..., Awaitable[Any]],
+    ) -> Callable[..., Coroutine[Any, Any, None]]:
+        @functools.wraps(func)
+        async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            try:
+                await func(self, *args, **kwargs)
+            except ApiException as exception:
+                raise HomeAssistantError(
+                    translation_domain=DOMAIN,
+                    translation_key="command_send_failed",
+                    translation_placeholders={"exception": str(exception)},
+                ) from exception
+            else:
+                if poll_after_sending:
+                    # As there are no updates from the websocket for this attribute,
+                    # we need to wait until the command is executed and then poll the API.
+                    await asyncio.sleep(EXECUTION_TIME_DELAY)
+                    await self.coordinator.async_request_refresh()
+
+        return wrapper
+
+    return decorator
+
+
 class AutomowerBaseEntity(CoordinatorEntity[AutomowerDataUpdateCoordinator]):
     """Defining the Automower base Entity."""
 

homeassistant/components/husqvarna_automower/lawn_mower.py

@@ -1,12 +1,8 @@
 """Husqvarna Automower lawn mower entity."""
 
-from collections.abc import Awaitable, Callable, Coroutine
 from datetime import timedelta
-import functools
 import logging
-from typing import Any
 
-from aioautomower.exceptions import ApiException
 from aioautomower.model import MowerActivities, MowerStates
 import voluptuous as vol
 
@@ -16,14 +12,12 @@ from homeassistant.components.lawn_mower import (
     LawnMowerEntityFeature,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import HomeAssistantError
 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
+from .entity import AutomowerAvailableEntity, handle_sending_exception
 
 DOCKED_ACTIVITIES = (MowerActivities.PARKED_IN_CS, MowerActivities.CHARGING)
 MOWING_ACTIVITIES = (
@@ -49,25 +43,6 @@ OVERRIDE_MODES = [MOW, PARK]
 _LOGGER = logging.getLogger(__name__)
 
 
-def handle_sending_exception(
-    func: Callable[..., Awaitable[Any]],
-) -> Callable[..., Coroutine[Any, Any, None]]:
-    """Handle exceptions while sending a command."""
-
-    @functools.wraps(func)
-    async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
-        try:
-            return await func(self, *args, **kwargs)
-        except ApiException as exception:
-            raise HomeAssistantError(
-                translation_domain=DOMAIN,
-                translation_key="command_send_failed",
-                translation_placeholders={"exception": str(exception)},
-            ) from exception
-
-    return wrapper
-
-
 async def async_setup_entry(
     hass: HomeAssistant,
     entry: AutomowerConfigEntry,
@@ -123,22 +98,22 @@ class AutomowerLawnMowerEntity(AutomowerAvailableEntity, LawnMowerEntity):
             return LawnMowerActivity.DOCKED
         return LawnMowerActivity.ERROR
 
-    @handle_sending_exception
+    @handle_sending_exception()
     async def async_start_mowing(self) -> None:
         """Resume schedule."""
         await self.coordinator.api.commands.resume_schedule(self.mower_id)
 
-    @handle_sending_exception
+    @handle_sending_exception()
     async def async_pause(self) -> None:
         """Pauses the mower."""
         await self.coordinator.api.commands.pause_mowing(self.mower_id)
 
-    @handle_sending_exception
+    @handle_sending_exception()
     async def async_dock(self) -> None:
         """Parks the mower until next schedule."""
         await self.coordinator.api.commands.park_until_next_schedule(self.mower_id)
 
-    @handle_sending_exception
+    @handle_sending_exception()
     async def async_override_schedule(
         self, override_mode: str, duration: timedelta
     ) -> None:

homeassistant/components/husqvarna_automower/number.py

@@ -1,26 +1,22 @@
 """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, WorkArea
 from aioautomower.session import AutomowerSession
 
 from homeassistant.components.number import NumberEntity, NumberEntityDescription
 from homeassistant.const import PERCENTAGE, EntityCategory, Platform
 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 . import AutomowerConfigEntry
-from .const import EXECUTION_TIME_DELAY
 from .coordinator import AutomowerDataUpdateCoordinator
-from .entity import AutomowerControlEntity
+from .entity import AutomowerControlEntity, handle_sending_exception
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -160,16 +156,12 @@ class AutomowerNumberEntity(AutomowerControlEntity, NumberEntity):
         """Return the state of the number."""
         return self.entity_description.value_fn(self.mower_attributes)
 
+    @handle_sending_exception()
     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.api, self.mower_id, value
-            )
-        except ApiException as exception:
-            raise HomeAssistantError(
-                f"Command couldn't be sent to the command queue: {exception}"
-            ) from exception
+        await self.entity_description.set_value_fn(
+            self.coordinator.api, self.mower_id, value
+        )
 
 
 class AutomowerWorkAreaNumberEntity(AutomowerControlEntity, NumberEntity):
@@ -208,21 +200,12 @@ class AutomowerWorkAreaNumberEntity(AutomowerControlEntity, NumberEntity):
         """Return the state of the number."""
         return self.entity_description.value_fn(self.work_area)
 
+    @handle_sending_exception(poll_after_sending=True)
     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
-        else:
-            # 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(EXECUTION_TIME_DELAY)
-            await self.coordinator.async_request_refresh()
+        await self.entity_description.set_value_fn(
+            self.coordinator, self.mower_id, value, self.work_area_id
+        )
 
 
 @callback

homeassistant/components/husqvarna_automower/select.py

@@ -3,18 +3,16 @@
 import logging
 from typing import cast
 
-from aioautomower.exceptions import ApiException
 from aioautomower.model import HeadlightModes
 
 from homeassistant.components.select import SelectEntity
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import AutomowerConfigEntry
 from .coordinator import AutomowerDataUpdateCoordinator
-from .entity import AutomowerControlEntity
+from .entity import AutomowerControlEntity, handle_sending_exception
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -64,13 +62,9 @@ class AutomowerSelectEntity(AutomowerControlEntity, SelectEntity):
             HeadlightModes, self.mower_attributes.settings.headlight.mode
         ).lower()
 
+    @handle_sending_exception()
     async def async_select_option(self, option: str) -> None:
         """Change the selected option."""
-        try:
-            await self.coordinator.api.commands.set_headlight_mode(
-                self.mower_id, cast(HeadlightModes, option.upper())
-            )
-        except ApiException as exception:
-            raise HomeAssistantError(
-                f"Command couldn't be sent to the command queue: {exception}"
-            ) from exception
+        await self.coordinator.api.commands.set_headlight_mode(
+            self.mower_id, cast(HeadlightModes, option.upper())
+        )

homeassistant/components/husqvarna_automower/switch.py

@@ -1,23 +1,19 @@
 """Creates a switch entity for the mower."""
 
-import asyncio
 import logging
 from typing import TYPE_CHECKING, Any
 
-from aioautomower.exceptions import ApiException
 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.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import AutomowerConfigEntry
-from .const import EXECUTION_TIME_DELAY
 from .coordinator import AutomowerDataUpdateCoordinator
-from .entity import AutomowerControlEntity
+from .entity import AutomowerControlEntity, handle_sending_exception
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -67,23 +63,15 @@ class AutomowerScheduleSwitchEntity(AutomowerControlEntity, SwitchEntity):
         """Return the state of the switch."""
         return self.mower_attributes.mower.mode != MowerModes.HOME
 
+    @handle_sending_exception()
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn the entity off."""
-        try:
-            await self.coordinator.api.commands.park_until_further_notice(self.mower_id)
-        except ApiException as exception:
-            raise HomeAssistantError(
-                f"Command couldn't be sent to the command queue: {exception}"
-            ) from exception
+        await self.coordinator.api.commands.park_until_further_notice(self.mower_id)
 
+    @handle_sending_exception()
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn the entity on."""
-        try:
-            await self.coordinator.api.commands.resume_schedule(self.mower_id)
-        except ApiException as exception:
-            raise HomeAssistantError(
-                f"Command couldn't be sent to the command queue: {exception}"
-            ) from exception
+        await self.coordinator.api.commands.resume_schedule(self.mower_id)
 
 
 class AutomowerStayOutZoneSwitchEntity(AutomowerControlEntity, SwitchEntity):
@@ -128,37 +116,19 @@ class AutomowerStayOutZoneSwitchEntity(AutomowerControlEntity, SwitchEntity):
         """Return True if the device is available and the zones are not `dirty`."""
         return super().available and not self.stay_out_zones.dirty
 
+    @handle_sending_exception(poll_after_sending=True)
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn the switch off."""
-        try:
-            await self.coordinator.api.commands.switch_stay_out_zone(
-                self.mower_id, self.stay_out_zone_uid, False
-            )
-        except ApiException as exception:
-            raise HomeAssistantError(
-                f"Command couldn't be sent to the command queue: {exception}"
-            ) from exception
-        else:
-            # As there are no updates from the websocket regarding stay out zone changes,
-            # we need to wait until the command is executed and then poll the API.
-            await asyncio.sleep(EXECUTION_TIME_DELAY)
-            await self.coordinator.async_request_refresh()
+        await self.coordinator.api.commands.switch_stay_out_zone(
+            self.mower_id, self.stay_out_zone_uid, False
+        )
 
+    @handle_sending_exception(poll_after_sending=True)
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn the switch on."""
-        try:
-            await self.coordinator.api.commands.switch_stay_out_zone(
-                self.mower_id, self.stay_out_zone_uid, True
-            )
-        except ApiException as exception:
-            raise HomeAssistantError(
-                f"Command couldn't be sent to the command queue: {exception}"
-            ) from exception
-        else:
-            # As there are no updates from the websocket regarding stay out zone changes,
-            # we need to wait until the command is executed and then poll the API.
-            await asyncio.sleep(EXECUTION_TIME_DELAY)
-            await self.coordinator.async_request_refresh()
+        await self.coordinator.api.commands.switch_stay_out_zone(
+            self.mower_id, self.stay_out_zone_uid, True
+        )
 
 
 @callback

tests/components/husqvarna_automower/test_number.py

@@ -41,7 +41,7 @@ async def test_number_commands(
     mocked_method.side_effect = ApiException("Test error")
     with pytest.raises(
         HomeAssistantError,
-        match="Command couldn't be sent to the command queue: Test error",
+        match="Failed to send command: Test error",
     ):
         await hass.services.async_call(
             domain="number",
@@ -85,7 +85,7 @@ async def test_number_workarea_commands(
     mocked_method.side_effect = ApiException("Test error")
     with pytest.raises(
         HomeAssistantError,
-        match="Command couldn't be sent to the command queue: Test error",
+        match="Failed to send command: Test error",
     ):
         await hass.services.async_call(
             domain="number",

tests/components/husqvarna_automower/test_select.py

@@ -88,7 +88,7 @@ async def test_select_commands(
     mocked_method.side_effect = ApiException("Test error")
     with pytest.raises(
         HomeAssistantError,
-        match="Command couldn't be sent to the command queue: Test error",
+        match="Failed to send command: Test error",
     ):
         await hass.services.async_call(
             domain="select",

tests/components/husqvarna_automower/test_switch.py

@@ -83,7 +83,7 @@ async def test_switch_commands(
     mocked_method.side_effect = ApiException("Test error")
     with pytest.raises(
         HomeAssistantError,
-        match="Command couldn't be sent to the command queue: Test error",
+        match="Failed to send command: Test error",
     ):
         await hass.services.async_call(
             domain="switch",
@@ -134,7 +134,7 @@ async def test_stay_out_zone_switch_commands(
     mocked_method.side_effect = ApiException("Test error")
     with pytest.raises(
         HomeAssistantError,
-        match="Command couldn't be sent to the command queue: Test error",
+        match="Failed to send command: Test error",
     ):
         await hass.services.async_call(
             domain="switch",