Store runtime data inside the config entry in Tautulli (#119552)

This commit is contained in:
Robert Hillis 2024-06-13 02:57:28 -04:00 committed by GitHub
parent 610f21c4a6
commit cad6163162
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 23 additions and 18 deletions

View file

@ -16,9 +16,10 @@ from .const import DEFAULT_NAME, DOMAIN
from .coordinator import TautulliDataUpdateCoordinator from .coordinator import TautulliDataUpdateCoordinator
PLATFORMS = [Platform.SENSOR] PLATFORMS = [Platform.SENSOR]
type TautulliConfigEntry = ConfigEntry[TautulliDataUpdateCoordinator]
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_setup_entry(hass: HomeAssistant, entry: TautulliConfigEntry) -> bool:
"""Set up Tautulli from a config entry.""" """Set up Tautulli from a config entry."""
host_configuration = PyTautulliHostConfiguration( host_configuration = PyTautulliHostConfiguration(
api_token=entry.data[CONF_API_KEY], api_token=entry.data[CONF_API_KEY],
@ -29,19 +30,18 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
host_configuration=host_configuration, host_configuration=host_configuration,
session=async_get_clientsession(hass, entry.data[CONF_VERIFY_SSL]), session=async_get_clientsession(hass, entry.data[CONF_VERIFY_SSL]),
) )
coordinator = TautulliDataUpdateCoordinator(hass, host_configuration, api_client) entry.runtime_data = TautulliDataUpdateCoordinator(
await coordinator.async_config_entry_first_refresh() hass, host_configuration, api_client
hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator )
await entry.runtime_data.async_config_entry_first_refresh()
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
return True return True
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, entry: TautulliConfigEntry) -> bool:
"""Unload a config entry.""" """Unload a config entry."""
if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
hass.data[DOMAIN].pop(entry.entry_id)
return unload_ok
class TautulliEntity(CoordinatorEntity[TautulliDataUpdateCoordinator]): class TautulliEntity(CoordinatorEntity[TautulliDataUpdateCoordinator]):

View file

@ -4,6 +4,7 @@ from __future__ import annotations
import asyncio import asyncio
from datetime import timedelta from datetime import timedelta
from typing import TYPE_CHECKING
from pytautulli import ( from pytautulli import (
PyTautulli, PyTautulli,
@ -17,18 +18,20 @@ from pytautulli.exceptions import (
) )
from pytautulli.models.host_configuration import PyTautulliHostConfiguration from pytautulli.models.host_configuration import PyTautulliHostConfiguration
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryAuthFailed from homeassistant.exceptions import ConfigEntryAuthFailed
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
from .const import DOMAIN, LOGGER from .const import DOMAIN, LOGGER
if TYPE_CHECKING:
from . import TautulliConfigEntry
class TautulliDataUpdateCoordinator(DataUpdateCoordinator[None]): class TautulliDataUpdateCoordinator(DataUpdateCoordinator[None]):
"""Data update coordinator for the Tautulli integration.""" """Data update coordinator for the Tautulli integration."""
config_entry: ConfigEntry config_entry: TautulliConfigEntry
def __init__( def __init__(
self, self,

View file

@ -19,14 +19,14 @@ from homeassistant.components.sensor import (
SensorEntityDescription, SensorEntityDescription,
SensorStateClass, SensorStateClass,
) )
from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry from homeassistant.config_entries import SOURCE_IMPORT
from homeassistant.const import PERCENTAGE, EntityCategory, UnitOfInformation from homeassistant.const import PERCENTAGE, EntityCategory, UnitOfInformation
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity import EntityDescription from homeassistant.helpers.entity import EntityDescription
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, StateType from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, StateType
from . import TautulliEntity from . import TautulliConfigEntry, TautulliEntity
from .const import ATTR_TOP_USER, DOMAIN from .const import ATTR_TOP_USER, DOMAIN
from .coordinator import TautulliDataUpdateCoordinator from .coordinator import TautulliDataUpdateCoordinator
@ -210,26 +210,28 @@ async def async_setup_platform(
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback hass: HomeAssistant,
entry: TautulliConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up Tautulli sensor.""" """Set up Tautulli sensor."""
coordinator: TautulliDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] data = entry.runtime_data
entities: list[TautulliSensor | TautulliSessionSensor] = [ entities: list[TautulliSensor | TautulliSessionSensor] = [
TautulliSensor( TautulliSensor(
coordinator, data,
description, description,
) )
for description in SENSOR_TYPES for description in SENSOR_TYPES
] ]
if coordinator.users: if data.users:
entities.extend( entities.extend(
TautulliSessionSensor( TautulliSessionSensor(
coordinator, data,
description, description,
user, user,
) )
for description in SESSION_SENSOR_TYPES for description in SESSION_SENSOR_TYPES
for user in coordinator.users for user in data.users
if user.username != "Local" if user.username != "Local"
) )
async_add_entities(entities) async_add_entities(entities)