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>
pull/116305/head^2
Mr. Bubbles 2024-05-05 17:02:28 +02:00 committed by GitHub
parent ffe6b9b6f0
commit b53081dc51
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 142 additions and 118 deletions

View File

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

View File

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

View File

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

View File

@ -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 aiohttp import ClientResponseError
from typing import TYPE_CHECKING, cast
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])
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 @@ class HabitipySensor(SensorEntity):
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 @@ 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:
all_received_tasks = self._updater.tasks
for element in self._task_type.path:
received_tasks = all_received_tasks[element]
attrs = {}
attrs = {}
# Map tasks to TASKS_MAP
for received_task in received_tasks:
# 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):

View File

@ -59,6 +59,11 @@
}
}
},
"exceptions": {
"setup_rate_limit_exception": {
"message": "Currently rate limited, try again later"
}
},
"services": {
"api_call": {
"name": "API name",

View File

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