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 <[email protected]>

* Update homeassistant/components/habitica/sensor.py

Co-authored-by: Joost Lekkerkerker <[email protected]>

* remove auth exception

* fixes

---------

Co-authored-by: Joost Lekkerkerker <[email protected]>

.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 @@
  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)

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)