From fe3027f7de279d1708d4cefb02503f9bf0facb3d Mon Sep 17 00:00:00 2001
From: Thomas55555 <59625598+Thomas55555@users.noreply.github.com>
Date: Mon, 24 Jun 2024 08:16:26 +0200
Subject: [PATCH] Adjust base entities in Husqvarna Automower (#120258)

* adjust base entities

* Adjust docstrings
---
 .../components/husqvarna_automower/button.py  |  4 +--
 .../components/husqvarna_automower/entity.py  | 33 +++++++++++++++++--
 .../husqvarna_automower/lawn_mower.py         |  4 +--
 .../components/husqvarna_automower/switch.py  | 31 +----------------
 4 files changed, 35 insertions(+), 37 deletions(-)

diff --git a/homeassistant/components/husqvarna_automower/button.py b/homeassistant/components/husqvarna_automower/button.py
index 60c05b92a31..a9747108393 100644
--- a/homeassistant/components/husqvarna_automower/button.py
+++ b/homeassistant/components/husqvarna_automower/button.py
@@ -12,7 +12,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from . import AutomowerConfigEntry
 from .const import DOMAIN
 from .coordinator import AutomowerDataUpdateCoordinator
-from .entity import AutomowerControlEntity
+from .entity import AutomowerAvailableEntity
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -29,7 +29,7 @@ async def async_setup_entry(
     )
 
 
-class AutomowerButtonEntity(AutomowerControlEntity, ButtonEntity):
+class AutomowerButtonEntity(AutomowerAvailableEntity, ButtonEntity):
     """Defining the AutomowerButtonEntity."""
 
     _attr_translation_key = "confirm_error"
diff --git a/homeassistant/components/husqvarna_automower/entity.py b/homeassistant/components/husqvarna_automower/entity.py
index 4d20d2d677b..80a936c2caf 100644
--- a/homeassistant/components/husqvarna_automower/entity.py
+++ b/homeassistant/components/husqvarna_automower/entity.py
@@ -2,7 +2,7 @@
 
 import logging
 
-from aioautomower.model import MowerAttributes
+from aioautomower.model import MowerActivities, MowerAttributes, MowerStates
 
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
@@ -12,6 +12,21 @@ from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
+ERROR_ACTIVITIES = (
+    MowerActivities.STOPPED_IN_GARDEN,
+    MowerActivities.UNKNOWN,
+    MowerActivities.NOT_APPLICABLE,
+)
+ERROR_STATES = [
+    MowerStates.FATAL_ERROR,
+    MowerStates.ERROR,
+    MowerStates.ERROR_AT_POWER_UP,
+    MowerStates.NOT_APPLICABLE,
+    MowerStates.UNKNOWN,
+    MowerStates.STOPPED,
+    MowerStates.OFF,
+]
+
 
 class AutomowerBaseEntity(CoordinatorEntity[AutomowerDataUpdateCoordinator]):
     """Defining the Automower base Entity."""
@@ -41,10 +56,22 @@ class AutomowerBaseEntity(CoordinatorEntity[AutomowerDataUpdateCoordinator]):
         return self.coordinator.data[self.mower_id]
 
 
-class AutomowerControlEntity(AutomowerBaseEntity):
-    """AutomowerControlEntity, for dynamic availability."""
+class AutomowerAvailableEntity(AutomowerBaseEntity):
+    """Replies available when the mower is connected."""
 
     @property
     def available(self) -> bool:
         """Return True if the device is available."""
         return super().available and self.mower_attributes.metadata.connected
+
+
+class AutomowerControlEntity(AutomowerAvailableEntity):
+    """Replies available when the mower is connected and not in error state."""
+
+    @property
+    def available(self) -> bool:
+        """Return True if the device is available."""
+        return super().available and (
+            self.mower_attributes.mower.state not in ERROR_STATES
+            or self.mower_attributes.mower.activity not in ERROR_ACTIVITIES
+        )
diff --git a/homeassistant/components/husqvarna_automower/lawn_mower.py b/homeassistant/components/husqvarna_automower/lawn_mower.py
index c0b566a7f66..e59d9e635e9 100644
--- a/homeassistant/components/husqvarna_automower/lawn_mower.py
+++ b/homeassistant/components/husqvarna_automower/lawn_mower.py
@@ -23,7 +23,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from . import AutomowerConfigEntry
 from .const import DOMAIN
 from .coordinator import AutomowerDataUpdateCoordinator
-from .entity import AutomowerControlEntity
+from .entity import AutomowerAvailableEntity
 
 DOCKED_ACTIVITIES = (MowerActivities.PARKED_IN_CS, MowerActivities.CHARGING)
 MOWING_ACTIVITIES = (
@@ -94,7 +94,7 @@ async def async_setup_entry(
     )
 
 
-class AutomowerLawnMowerEntity(AutomowerControlEntity, LawnMowerEntity):
+class AutomowerLawnMowerEntity(AutomowerAvailableEntity, LawnMowerEntity):
     """Defining each mower Entity."""
 
     _attr_name = None
diff --git a/homeassistant/components/husqvarna_automower/switch.py b/homeassistant/components/husqvarna_automower/switch.py
index a856e9c9050..8a450b8e81a 100644
--- a/homeassistant/components/husqvarna_automower/switch.py
+++ b/homeassistant/components/husqvarna_automower/switch.py
@@ -5,13 +5,7 @@ import logging
 from typing import TYPE_CHECKING, Any
 
 from aioautomower.exceptions import ApiException
-from aioautomower.model import (
-    MowerActivities,
-    MowerModes,
-    MowerStates,
-    StayOutZones,
-    Zone,
-)
+from aioautomower.model import MowerModes, StayOutZones, Zone
 
 from homeassistant.components.switch import SwitchEntity
 from homeassistant.const import Platform
@@ -27,21 +21,6 @@ from .entity import AutomowerControlEntity
 
 _LOGGER = logging.getLogger(__name__)
 
-ERROR_ACTIVITIES = (
-    MowerActivities.STOPPED_IN_GARDEN,
-    MowerActivities.UNKNOWN,
-    MowerActivities.NOT_APPLICABLE,
-)
-ERROR_STATES = [
-    MowerStates.FATAL_ERROR,
-    MowerStates.ERROR,
-    MowerStates.ERROR_AT_POWER_UP,
-    MowerStates.NOT_APPLICABLE,
-    MowerStates.UNKNOWN,
-    MowerStates.STOPPED,
-    MowerStates.OFF,
-]
-
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -88,14 +67,6 @@ class AutomowerScheduleSwitchEntity(AutomowerControlEntity, SwitchEntity):
         """Return the state of the switch."""
         return self.mower_attributes.mower.mode != MowerModes.HOME
 
-    @property
-    def available(self) -> bool:
-        """Return True if the device is available."""
-        return super().available and (
-            self.mower_attributes.mower.state not in ERROR_STATES
-            or self.mower_attributes.mower.activity not in ERROR_ACTIVITIES
-        )
-
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn the entity off."""
         try: