Refactor Habitica button and switch functions to use habiticalib instance directly (#149602)

homeassistant/components/habitica/button.py

@@ -7,15 +7,7 @@ from dataclasses import dataclass
 from enum import StrEnum
 from typing import Any
 
-from aiohttp import ClientError
-from habiticalib import (
-    HabiticaClass,
-    HabiticaException,
-    NotAuthorizedError,
-    Skill,
-    TaskType,
-    TooManyRequestsError,
-)
+from habiticalib import Habitica, HabiticaClass, Skill, TaskType
 
 from homeassistant.components.button import (
     DOMAIN as BUTTON_DOMAIN,
@@ -23,16 +15,11 @@ from homeassistant.components.button import (
     ButtonEntityDescription,
 )
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
 from .const import ASSETS_URL, DOMAIN
-from .coordinator import (
-    HabiticaConfigEntry,
-    HabiticaData,
-    HabiticaDataUpdateCoordinator,
-)
+from .coordinator import HabiticaConfigEntry, HabiticaData
 from .entity import HabiticaBase
 
 PARALLEL_UPDATES = 1
@@ -42,7 +29,7 @@ PARALLEL_UPDATES = 1
 class HabiticaButtonEntityDescription(ButtonEntityDescription):
     """Describes Habitica button entity."""
 
-    press_fn: Callable[[HabiticaDataUpdateCoordinator], Any]
+    press_fn: Callable[[Habitica], Any]
     available_fn: Callable[[HabiticaData], bool]
     class_needed: HabiticaClass | None = None
     entity_picture: str | None = None
@@ -73,13 +60,13 @@ BUTTON_DESCRIPTIONS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.RUN_CRON,
         translation_key=HabiticaButtonEntity.RUN_CRON,
-        press_fn=lambda coordinator: coordinator.habitica.run_cron(),
+        press_fn=lambda habitica: habitica.run_cron(),
         available_fn=lambda data: data.user.needsCron is True,
     ),
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.BUY_HEALTH_POTION,
         translation_key=HabiticaButtonEntity.BUY_HEALTH_POTION,
-        press_fn=lambda coordinator: coordinator.habitica.buy_health_potion(),
+        press_fn=lambda habitica: habitica.buy_health_potion(),
         available_fn=(
             lambda data: (data.user.stats.gp or 0) >= 25
             and (data.user.stats.hp or 0) < 50
@@ -89,7 +76,7 @@ BUTTON_DESCRIPTIONS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.ALLOCATE_ALL_STAT_POINTS,
         translation_key=HabiticaButtonEntity.ALLOCATE_ALL_STAT_POINTS,
-        press_fn=lambda coordinator: coordinator.habitica.allocate_stat_points(),
+        press_fn=lambda habitica: habitica.allocate_stat_points(),
         available_fn=(
             lambda data: data.user.preferences.automaticAllocation is True
             and (data.user.stats.points or 0) > 0
@@ -98,7 +85,7 @@ BUTTON_DESCRIPTIONS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.REVIVE,
         translation_key=HabiticaButtonEntity.REVIVE,
-        press_fn=lambda coordinator: coordinator.habitica.revive(),
+        press_fn=lambda habitica: habitica.revive(),
         available_fn=lambda data: data.user.stats.hp == 0,
     ),
 )
@@ -108,9 +95,7 @@ CLASS_SKILLS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.MPHEAL,
         translation_key=HabiticaButtonEntity.MPHEAL,
-        press_fn=(
-            lambda coordinator: coordinator.habitica.cast_skill(Skill.ETHEREAL_SURGE)
-        ),
+        press_fn=lambda habitica: habitica.cast_skill(Skill.ETHEREAL_SURGE),
         available_fn=(
             lambda data: (data.user.stats.lvl or 0) >= 12
             and (data.user.stats.mp or 0) >= 30
@@ -121,7 +106,7 @@ CLASS_SKILLS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.EARTH,
         translation_key=HabiticaButtonEntity.EARTH,
-        press_fn=lambda coordinator: coordinator.habitica.cast_skill(Skill.EARTHQUAKE),
+        press_fn=lambda habitica: habitica.cast_skill(Skill.EARTHQUAKE),
         available_fn=(
             lambda data: (data.user.stats.lvl or 0) >= 13
             and (data.user.stats.mp or 0) >= 35
@@ -132,9 +117,7 @@ CLASS_SKILLS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.FROST,
         translation_key=HabiticaButtonEntity.FROST,
-        press_fn=(
-            lambda coordinator: coordinator.habitica.cast_skill(Skill.CHILLING_FROST)
-        ),
+        press_fn=lambda habitica: habitica.cast_skill(Skill.CHILLING_FROST),
         # chilling frost can only be cast once per day (streaks buff is false)
         available_fn=(
             lambda data: (data.user.stats.lvl or 0) >= 14
@@ -147,9 +130,7 @@ CLASS_SKILLS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.DEFENSIVE_STANCE,
         translation_key=HabiticaButtonEntity.DEFENSIVE_STANCE,
-        press_fn=(
-            lambda coordinator: coordinator.habitica.cast_skill(Skill.DEFENSIVE_STANCE)
-        ),
+        press_fn=lambda habitica: habitica.cast_skill(Skill.DEFENSIVE_STANCE),
         available_fn=(
             lambda data: (data.user.stats.lvl or 0) >= 12
             and (data.user.stats.mp or 0) >= 25
@@ -160,9 +141,7 @@ CLASS_SKILLS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.VALOROUS_PRESENCE,
         translation_key=HabiticaButtonEntity.VALOROUS_PRESENCE,
-        press_fn=(
-            lambda coordinator: coordinator.habitica.cast_skill(Skill.VALOROUS_PRESENCE)
-        ),
+        press_fn=lambda habitica: habitica.cast_skill(Skill.VALOROUS_PRESENCE),
         available_fn=(
             lambda data: (data.user.stats.lvl or 0) >= 13
             and (data.user.stats.mp or 0) >= 20
@@ -173,9 +152,7 @@ CLASS_SKILLS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.INTIMIDATE,
         translation_key=HabiticaButtonEntity.INTIMIDATE,
-        press_fn=(
-            lambda coordinator: coordinator.habitica.cast_skill(Skill.INTIMIDATING_GAZE)
-        ),
+        press_fn=lambda habitica: habitica.cast_skill(Skill.INTIMIDATING_GAZE),
         available_fn=(
             lambda data: (data.user.stats.lvl or 0) >= 14
             and (data.user.stats.mp or 0) >= 15
@@ -186,11 +163,7 @@ CLASS_SKILLS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.TOOLS_OF_TRADE,
         translation_key=HabiticaButtonEntity.TOOLS_OF_TRADE,
-        press_fn=(
-            lambda coordinator: coordinator.habitica.cast_skill(
-                Skill.TOOLS_OF_THE_TRADE
-            )
-        ),
+        press_fn=lambda habitica: habitica.cast_skill(Skill.TOOLS_OF_THE_TRADE),
         available_fn=(
             lambda data: (data.user.stats.lvl or 0) >= 13
             and (data.user.stats.mp or 0) >= 25
@@ -201,7 +174,7 @@ CLASS_SKILLS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.STEALTH,
         translation_key=HabiticaButtonEntity.STEALTH,
-        press_fn=lambda coordinator: coordinator.habitica.cast_skill(Skill.STEALTH),
+        press_fn=lambda habitica: habitica.cast_skill(Skill.STEALTH),
         # Stealth buffs stack and it can only be cast if the amount of
         # buffs is smaller than the amount of unfinished dailies
         available_fn=(
@@ -224,9 +197,7 @@ CLASS_SKILLS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.HEAL,
         translation_key=HabiticaButtonEntity.HEAL,
-        press_fn=(
-            lambda coordinator: coordinator.habitica.cast_skill(Skill.HEALING_LIGHT)
-        ),
+        press_fn=lambda habitica: habitica.cast_skill(Skill.HEALING_LIGHT),
         available_fn=(
             lambda data: (data.user.stats.lvl or 0) >= 11
             and (data.user.stats.mp or 0) >= 15
@@ -238,11 +209,7 @@ CLASS_SKILLS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.BRIGHTNESS,
         translation_key=HabiticaButtonEntity.BRIGHTNESS,
-        press_fn=(
-            lambda coordinator: coordinator.habitica.cast_skill(
-                Skill.SEARING_BRIGHTNESS
-            )
-        ),
+        press_fn=lambda habitica: habitica.cast_skill(Skill.SEARING_BRIGHTNESS),
         available_fn=(
             lambda data: (data.user.stats.lvl or 0) >= 12
             and (data.user.stats.mp or 0) >= 15
@@ -253,9 +220,7 @@ CLASS_SKILLS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.PROTECT_AURA,
         translation_key=HabiticaButtonEntity.PROTECT_AURA,
-        press_fn=(
-            lambda coordinator: coordinator.habitica.cast_skill(Skill.PROTECTIVE_AURA)
-        ),
+        press_fn=lambda habitica: habitica.cast_skill(Skill.PROTECTIVE_AURA),
         available_fn=(
             lambda data: (data.user.stats.lvl or 0) >= 13
             and (data.user.stats.mp or 0) >= 30
@@ -266,7 +231,7 @@ CLASS_SKILLS: tuple[HabiticaButtonEntityDescription, ...] = (
     HabiticaButtonEntityDescription(
         key=HabiticaButtonEntity.HEAL_ALL,
         translation_key=HabiticaButtonEntity.HEAL_ALL,
-        press_fn=lambda coordinator: coordinator.habitica.cast_skill(Skill.BLESSING),
+        press_fn=lambda habitica: habitica.cast_skill(Skill.BLESSING),
         available_fn=(
             lambda data: (data.user.stats.lvl or 0) >= 14
             and (data.user.stats.mp or 0) >= 25
@@ -332,33 +297,9 @@ class HabiticaButton(HabiticaBase, ButtonEntity):
 
     async def async_press(self) -> None:
         """Handle the button press."""
-        try:
-            await self.entity_description.press_fn(self.coordinator)
-        except TooManyRequestsError as e:
-            raise HomeAssistantError(
-                translation_domain=DOMAIN,
-                translation_key="setup_rate_limit_exception",
-                translation_placeholders={"retry_after": str(e.retry_after)},
-            ) from e
-        except NotAuthorizedError as e:
-            raise ServiceValidationError(
-                translation_domain=DOMAIN,
-                translation_key="service_call_unallowed",
-            ) from e
-        except HabiticaException as e:
-            raise HomeAssistantError(
-                translation_domain=DOMAIN,
-                translation_key="service_call_exception",
-                translation_placeholders={"reason": e.error.message},
-            ) from e
-        except ClientError as e:
-            raise HomeAssistantError(
-                translation_domain=DOMAIN,
-                translation_key="service_call_exception",
-                translation_placeholders={"reason": str(e)},
-            ) from e
-        else:
-            await self.coordinator.async_request_refresh()
+
+        await self.coordinator.execute(self.entity_description.press_fn)
+        await self.coordinator.async_request_refresh()
 
     @property
     def available(self) -> bool:

homeassistant/components/habitica/coordinator.py

@@ -28,6 +28,7 @@ from homeassistant.exceptions import (
     ConfigEntryAuthFailed,
     ConfigEntryNotReady,
     HomeAssistantError,
+    ServiceValidationError,
 )
 from homeassistant.helpers.debounce import Debouncer
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
@@ -130,19 +131,22 @@ class HabiticaDataUpdateCoordinator(DataUpdateCoordinator[HabiticaData]):
         else:
             return HabiticaData(user=user, tasks=tasks + completed_todos)
 
-    async def execute(
-        self, func: Callable[[HabiticaDataUpdateCoordinator], Any]
-    ) -> None:
+    async def execute(self, func: Callable[[Habitica], Any]) -> None:
         """Execute an API call."""
 
         try:
-            await func(self)
+            await func(self.habitica)
         except TooManyRequestsError as e:
             raise HomeAssistantError(
                 translation_domain=DOMAIN,
                 translation_key="setup_rate_limit_exception",
                 translation_placeholders={"retry_after": str(e.retry_after)},
             ) from e
+        except NotAuthorizedError as e:
+            raise ServiceValidationError(
+                translation_domain=DOMAIN,
+                translation_key="service_call_unallowed",
+            ) from e
         except HabiticaException as e:
             raise HomeAssistantError(
                 translation_domain=DOMAIN,

homeassistant/components/habitica/switch.py

@@ -7,6 +7,8 @@ from dataclasses import dataclass
 from enum import StrEnum
 from typing import Any
 
+from habiticalib import Habitica
+
 from homeassistant.components.switch import (
     SwitchDeviceClass,
     SwitchEntity,
@@ -15,11 +17,7 @@ from homeassistant.components.switch import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .coordinator import (
-    HabiticaConfigEntry,
-    HabiticaData,
-    HabiticaDataUpdateCoordinator,
-)
+from .coordinator import HabiticaConfigEntry, HabiticaData
 from .entity import HabiticaBase
 
 PARALLEL_UPDATES = 1
@@ -29,8 +27,8 @@ PARALLEL_UPDATES = 1
 class HabiticaSwitchEntityDescription(SwitchEntityDescription):
     """Describes Habitica switch entity."""
 
-    turn_on_fn: Callable[[HabiticaDataUpdateCoordinator], Any]
-    turn_off_fn: Callable[[HabiticaDataUpdateCoordinator], Any]
+    turn_on_fn: Callable[[Habitica], Any]
+    turn_off_fn: Callable[[Habitica], Any]
     is_on_fn: Callable[[HabiticaData], bool | None]
 
 
@@ -45,8 +43,8 @@ SWTICH_DESCRIPTIONS: tuple[HabiticaSwitchEntityDescription, ...] = (
         key=HabiticaSwitchEntity.SLEEP,
         translation_key=HabiticaSwitchEntity.SLEEP,
         device_class=SwitchDeviceClass.SWITCH,
-        turn_on_fn=lambda coordinator: coordinator.habitica.toggle_sleep(),
-        turn_off_fn=lambda coordinator: coordinator.habitica.toggle_sleep(),
+        turn_on_fn=lambda habitica: habitica.toggle_sleep(),
+        turn_off_fn=lambda habitica: habitica.toggle_sleep(),
         is_on_fn=lambda data: data.user.preferences.sleep,
     ),
 )