From b53081dc513a58ecc706e431bbddacf30ba17457 Mon Sep 17 00:00:00 2001 From: "Mr. Bubbles" Date: Sun, 5 May 2024 17:02:28 +0200 Subject: [PATCH] Add update coordinator for Habitica integration (#116427) * Add DataUpdateCoordinator and exception handling for service * remove unnecessary lines * revert changes to service * remove type check * store coordinator in config_entry * add exception translations * update HabiticaData * Update homeassistant/components/habitica/__init__.py Co-authored-by: Joost Lekkerkerker * Update homeassistant/components/habitica/sensor.py Co-authored-by: Joost Lekkerkerker * remove auth exception * fixes --------- Co-authored-by: Joost Lekkerkerker --- .coveragerc | 1 + homeassistant/components/habitica/__init__.py | 54 ++++--- .../components/habitica/coordinator.py | 56 +++++++ homeassistant/components/habitica/sensor.py | 138 +++++------------- .../components/habitica/strings.json | 5 + tests/components/habitica/test_init.py | 8 +- 6 files changed, 143 insertions(+), 119 deletions(-) create mode 100644 homeassistant/components/habitica/coordinator.py diff --git a/.coveragerc b/.coveragerc index 10dedd43e819d7..7986785d86e7a3 100644 --- a/.coveragerc +++ b/.coveragerc @@ -519,6 +519,7 @@ omit = homeassistant/components/guardian/util.py homeassistant/components/guardian/valve.py homeassistant/components/habitica/__init__.py + homeassistant/components/habitica/coordinator.py homeassistant/components/habitica/sensor.py homeassistant/components/harman_kardon_avr/media_player.py homeassistant/components/harmony/data.py diff --git a/homeassistant/components/habitica/__init__.py b/homeassistant/components/habitica/__init__.py index 34736116a26481..f5997b4a963060 100644 --- a/homeassistant/components/habitica/__init__.py +++ b/homeassistant/components/habitica/__init__.py @@ -1,7 +1,9 @@ """The habitica integration.""" +from http import HTTPStatus import logging +from aiohttp import ClientResponseError from habitipy.aio import HabitipyAsync import voluptuous as vol @@ -16,6 +18,7 @@ Platform, ) from homeassistant.core import HomeAssistant, ServiceCall +from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers import config_validation as cv from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.typing import ConfigType @@ -30,9 +33,12 @@ EVENT_API_CALL_SUCCESS, SERVICE_API_CALL, ) +from .coordinator import HabiticaDataUpdateCoordinator _LOGGER = logging.getLogger(__name__) +HabiticaConfigEntry = ConfigEntry[HabiticaDataUpdateCoordinator] + SENSORS_TYPES = ["name", "hp", "maxHealth", "mp", "maxMP", "exp", "toNextLevel", "lvl"] INSTANCE_SCHEMA = vol.All( @@ -104,7 +110,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: return True -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: HabiticaConfigEntry) -> bool: """Set up habitica from a config entry.""" class HAHabitipyAsync(HabitipyAsync): @@ -120,7 +126,7 @@ async def handle_api_call(call: ServiceCall) -> None: api = None for entry in entries: if entry.data[CONF_NAME] == name: - api = hass.data[DOMAIN].get(entry.entry_id) + api = entry.runtime_data.api break if api is None: _LOGGER.error("API_CALL: User '%s' not configured", name) @@ -139,24 +145,40 @@ async def handle_api_call(call: ServiceCall) -> None: EVENT_API_CALL_SUCCESS, {ATTR_NAME: name, ATTR_PATH: path, ATTR_DATA: data} ) - data = hass.data.setdefault(DOMAIN, {}) - config = entry.data websession = async_get_clientsession(hass) - url = config[CONF_URL] - username = config[CONF_API_USER] - password = config[CONF_API_KEY] - name = config.get(CONF_NAME) - config_dict = {"url": url, "login": username, "password": password} - api = HAHabitipyAsync(config_dict) - user = await api.user.get() - if name is None: + + url = entry.data[CONF_URL] + username = entry.data[CONF_API_USER] + password = entry.data[CONF_API_KEY] + + api = HAHabitipyAsync( + { + "url": url, + "login": username, + "password": password, + } + ) + try: + user = await api.user.get(userFields="profile") + except ClientResponseError as e: + if e.status == HTTPStatus.TOO_MANY_REQUESTS: + raise ConfigEntryNotReady( + translation_domain=DOMAIN, + translation_key="setup_rate_limit_exception", + ) from e + raise ConfigEntryNotReady(e) from e + + if not entry.data.get(CONF_NAME): name = user["profile"]["name"] hass.config_entries.async_update_entry( entry, data={**entry.data, CONF_NAME: name}, ) - data[entry.entry_id] = api + coordinator = HabiticaDataUpdateCoordinator(hass, api) + await coordinator.async_config_entry_first_refresh() + + entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) if not hass.services.has_service(DOMAIN, SERVICE_API_CALL): @@ -169,10 +191,6 @@ async def handle_api_call(call: ServiceCall) -> None: async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload a config entry.""" - unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - if len(hass.config_entries.async_entries(DOMAIN)) == 1: hass.services.async_remove(DOMAIN, SERVICE_API_CALL) - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/habitica/coordinator.py b/homeassistant/components/habitica/coordinator.py new file mode 100644 index 00000000000000..385652f710a9a3 --- /dev/null +++ b/homeassistant/components/habitica/coordinator.py @@ -0,0 +1,56 @@ +"""DataUpdateCoordinator for the Habitica integration.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import timedelta +import logging +from typing import Any + +from aiohttp import ClientResponseError +from habitipy.aio import HabitipyAsync + +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import HomeAssistant +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed + +from .const import DOMAIN + +_LOGGER = logging.getLogger(__name__) + + +@dataclass +class HabiticaData: + """Coordinator data class.""" + + user: dict[str, Any] + tasks: list[dict] + + +class HabiticaDataUpdateCoordinator(DataUpdateCoordinator[HabiticaData]): + """Habitica Data Update Coordinator.""" + + config_entry: ConfigEntry + + def __init__(self, hass: HomeAssistant, habitipy: HabitipyAsync) -> None: + """Initialize the Habitica data coordinator.""" + super().__init__( + hass, + _LOGGER, + name=DOMAIN, + update_interval=timedelta(seconds=30), + ) + self.api = habitipy + + async def _async_update_data(self) -> HabiticaData: + user_fields = set(self.async_contexts()) + + try: + user_response = await self.api.user.get(userFields=",".join(user_fields)) + tasks_response = [] + for task_type in ("todos", "dailys", "habits", "rewards"): + tasks_response.extend(await self.api.tasks.user.get(type=task_type)) + except ClientResponseError as error: + raise UpdateFailed(f"Error communicating with API: {error}") from error + + return HabiticaData(user=user_response, tasks=tasks_response) diff --git a/homeassistant/components/habitica/sensor.py b/homeassistant/components/habitica/sensor.py index 7ced7cbf192d49..5073c31d350f81 100644 --- a/homeassistant/components/habitica/sensor.py +++ b/homeassistant/components/habitica/sensor.py @@ -4,13 +4,9 @@ from collections import namedtuple from dataclasses import dataclass -from datetime import timedelta from enum import StrEnum -from http import HTTPStatus import logging -from typing import TYPE_CHECKING, Any - -from aiohttp import ClientResponseError +from typing import TYPE_CHECKING, cast from homeassistant.components.sensor import ( SensorDeviceClass, @@ -22,14 +18,15 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback -from homeassistant.util import Throttle +from homeassistant.helpers.typing import StateType +from homeassistant.helpers.update_coordinator import CoordinatorEntity +from . import HabiticaConfigEntry from .const import DOMAIN, MANUFACTURER, NAME +from .coordinator import HabiticaDataUpdateCoordinator _LOGGER = logging.getLogger(__name__) -MIN_TIME_BETWEEN_UPDATES = timedelta(minutes=15) - @dataclass(kw_only=True, frozen=True) class HabitipySensorEntityDescription(SensorEntityDescription): @@ -122,14 +119,14 @@ class HabitipySensorEntity(StrEnum): SensorType = namedtuple("SensorType", ["name", "icon", "unit", "path"]) TASKS_TYPES = { "habits": SensorType( - "Habits", "mdi:clipboard-list-outline", "n_of_tasks", ["habits"] + "Habits", "mdi:clipboard-list-outline", "n_of_tasks", ["habit"] ), "dailys": SensorType( - "Dailys", "mdi:clipboard-list-outline", "n_of_tasks", ["dailys"] + "Dailys", "mdi:clipboard-list-outline", "n_of_tasks", ["daily"] ), - "todos": SensorType("TODOs", "mdi:clipboard-list-outline", "n_of_tasks", ["todos"]), + "todos": SensorType("TODOs", "mdi:clipboard-list-outline", "n_of_tasks", ["todo"]), "rewards": SensorType( - "Rewards", "mdi:clipboard-list-outline", "n_of_tasks", ["rewards"] + "Rewards", "mdi:clipboard-list-outline", "n_of_tasks", ["reward"] ), } @@ -163,79 +160,26 @@ class HabitipySensorEntity(StrEnum): async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: HabiticaConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up the habitica sensors.""" name = config_entry.data[CONF_NAME] - sensor_data = HabitipyData(hass.data[DOMAIN][config_entry.entry_id]) - await sensor_data.update() + coordinator = config_entry.runtime_data entities: list[SensorEntity] = [ - HabitipySensor(sensor_data, description, config_entry) + HabitipySensor(coordinator, description, config_entry) for description in SENSOR_DESCRIPTIONS.values() ] entities.extend( - HabitipyTaskSensor(name, task_type, sensor_data, config_entry) + HabitipyTaskSensor(name, task_type, coordinator, config_entry) for task_type in TASKS_TYPES ) async_add_entities(entities, True) -class HabitipyData: - """Habitica API user data cache.""" - - tasks: dict[str, Any] - - def __init__(self, api) -> None: - """Habitica API user data cache.""" - self.api = api - self.data = None - self.tasks = {} - - @Throttle(MIN_TIME_BETWEEN_UPDATES) - async def update(self): - """Get a new fix from Habitica servers.""" - try: - self.data = await self.api.user.get() - except ClientResponseError as error: - if error.status == HTTPStatus.TOO_MANY_REQUESTS: - _LOGGER.warning( - ( - "Sensor data update for %s has too many API requests;" - " Skipping the update" - ), - DOMAIN, - ) - else: - _LOGGER.error( - "Count not update sensor data for %s (%s)", - DOMAIN, - error, - ) - - for task_type in TASKS_TYPES: - try: - self.tasks[task_type] = await self.api.tasks.user.get(type=task_type) - except ClientResponseError as error: - if error.status == HTTPStatus.TOO_MANY_REQUESTS: - _LOGGER.warning( - ( - "Sensor data update for %s has too many API requests;" - " Skipping the update" - ), - DOMAIN, - ) - else: - _LOGGER.error( - "Count not update sensor data for %s (%s)", - DOMAIN, - error, - ) - - -class HabitipySensor(SensorEntity): +class HabitipySensor(CoordinatorEntity[HabiticaDataUpdateCoordinator], SensorEntity): """A generic Habitica sensor.""" _attr_has_entity_name = True @@ -243,15 +187,14 @@ class HabitipySensor(SensorEntity): def __init__( self, - coordinator, + coordinator: HabiticaDataUpdateCoordinator, entity_description: HabitipySensorEntityDescription, entry: ConfigEntry, ) -> None: """Initialize a generic Habitica sensor.""" - super().__init__() + super().__init__(coordinator, context=entity_description.value_path[0]) if TYPE_CHECKING: assert entry.unique_id - self.coordinator = coordinator self.entity_description = entity_description self._attr_unique_id = f"{entry.unique_id}_{entity_description.key}" self._attr_device_info = DeviceInfo( @@ -263,25 +206,27 @@ def __init__( identifiers={(DOMAIN, entry.unique_id)}, ) - async def async_update(self) -> None: - """Update Sensor state.""" - await self.coordinator.update() - data = self.coordinator.data + @property + def native_value(self) -> StateType: + """Return the state of the device.""" + data = self.coordinator.data.user for element in self.entity_description.value_path: data = data[element] - self._attr_native_value = data + return cast(StateType, data) -class HabitipyTaskSensor(SensorEntity): +class HabitipyTaskSensor( + CoordinatorEntity[HabiticaDataUpdateCoordinator], SensorEntity +): """A Habitica task sensor.""" - def __init__(self, name, task_name, updater, entry): + def __init__(self, name, task_name, coordinator, entry): """Initialize a generic Habitica task.""" + super().__init__(coordinator) self._name = name self._task_name = task_name self._task_type = TASKS_TYPES[task_name] self._state = None - self._updater = updater self._attr_unique_id = f"{entry.unique_id}_{task_name}" self._attr_device_info = DeviceInfo( entry_type=DeviceEntryType.SERVICE, @@ -292,14 +237,6 @@ def __init__(self, name, task_name, updater, entry): identifiers={(DOMAIN, entry.unique_id)}, ) - async def async_update(self) -> None: - """Update Condition and Forecast.""" - await self._updater.update() - all_tasks = self._updater.tasks - for element in self._task_type.path: - tasks_length = len(all_tasks[element]) - self._state = tasks_length - @property def icon(self): """Return the icon to use in the frontend, if any.""" @@ -313,26 +250,29 @@ def name(self): @property def native_value(self): """Return the state of the device.""" - return self._state + return len( + [ + task + for task in self.coordinator.data.tasks + if task.get("type") in self._task_type.path + ] + ) @property def extra_state_attributes(self): """Return the state attributes of all user tasks.""" - if self._updater.tasks is not None: - all_received_tasks = self._updater.tasks - for element in self._task_type.path: - received_tasks = all_received_tasks[element] - attrs = {} - - # Map tasks to TASKS_MAP - for received_task in received_tasks: + attrs = {} + + # Map tasks to TASKS_MAP + for received_task in self.coordinator.data.tasks: + if received_task.get("type") in self._task_type.path: task_id = received_task[TASKS_MAP_ID] task = {} for map_key, map_value in TASKS_MAP.items(): if value := received_task.get(map_value): task[map_key] = value attrs[task_id] = task - return attrs + return attrs @property def native_unit_of_measurement(self): diff --git a/homeassistant/components/habitica/strings.json b/homeassistant/components/habitica/strings.json index 6be2bd7ed09b71..6023aa2d228486 100644 --- a/homeassistant/components/habitica/strings.json +++ b/homeassistant/components/habitica/strings.json @@ -59,6 +59,11 @@ } } }, + "exceptions": { + "setup_rate_limit_exception": { + "message": "Currently rate limited, try again later" + } + }, "services": { "api_call": { "name": "API name", diff --git a/tests/components/habitica/test_init.py b/tests/components/habitica/test_init.py index 9168e29f2d506e..50c7e664cd4078 100644 --- a/tests/components/habitica/test_init.py +++ b/tests/components/habitica/test_init.py @@ -55,7 +55,7 @@ def common_requests(aioclient_mock): "api_user": "test-api-user", "profile": {"name": TEST_USER_NAME}, "stats": { - "class": "test-class", + "class": "warrior", "con": 1, "exp": 2, "gp": 3, @@ -78,7 +78,11 @@ def common_requests(aioclient_mock): f"https://habitica.com/api/v3/tasks/user?type={task_type}", json={ "data": [ - {"text": f"this is a mock {task_type} #{task}", "id": f"{task}"} + { + "text": f"this is a mock {task_type} #{task}", + "id": f"{task}", + "type": TASKS_TYPES[task_type].path[0], + } for task in range(n_tasks) ] },