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

* Update homeassistant/components/habitica/sensor.py

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

* remove auth exception

* fixes

---------

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

.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

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 @@ from homeassistant.const import (
     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 @@ from .const import (
     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 async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         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 async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             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]
+    url = entry.data[CONF_URL]
-    password = config[CONF_API_KEY]
+    username = entry.data[CONF_API_USER]
-    name = config.get(CONF_NAME)
+    password = entry.data[CONF_API_KEY]
-    config_dict = {"url": url, "login": username, "password": password}
+
-    api = HAHabitipyAsync(config_dict)
+    api = HAHabitipyAsync(
-    user = await api.user.get()
+        {
-    if name is None:
+            "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 async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 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)

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)

homeassistant/components/habitica/sensor.py

@@ -4,13 +4,9 @@ from __future__ import annotations
 
 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 typing import TYPE_CHECKING, cast
-
-from aiohttp import ClientResponseError
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
@@ -22,14 +18,15 @@ from homeassistant.const import CONF_NAME, CONF_URL
 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 @@ SENSOR_DESCRIPTIONS: dict[str, HabitipySensorEntityDescription] = {
 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 @@ TASKS_MAP = {
 
 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])
+    coordinator = config_entry.runtime_data
-    await sensor_data.update()
 
     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:
+class HabitipySensor(CoordinatorEntity[HabiticaDataUpdateCoordinator], SensorEntity):
-    """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):
     """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 @@ class HabitipySensor(SensorEntity):
             identifiers={(DOMAIN, entry.unique_id)},
         )
 
-    async def async_update(self) -> None:
+    @property
-        """Update Sensor state."""
+    def native_value(self) -> StateType:
-        await self.coordinator.update()
+        """Return the state of the device."""
-        data = self.coordinator.data
+        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 @@ class HabitipyTaskSensor(SensorEntity):
             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 @@ class HabitipyTaskSensor(SensorEntity):
     @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:
+        attrs = {}
-            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
+        # Map tasks to TASKS_MAP
-            for received_task in received_tasks:
+        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):

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",

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)
                 ]
             },