From cad616316205311492d836c748f1a32da1500fec Mon Sep 17 00:00:00 2001
From: Robert Hillis <tkdrob4390@yahoo.com>
Date: Thu, 13 Jun 2024 02:57:28 -0400
Subject: [PATCH] Store runtime data inside the config entry in Tautulli
 (#119552)

---
 homeassistant/components/tautulli/__init__.py  | 16 ++++++++--------
 .../components/tautulli/coordinator.py         |  7 +++++--
 homeassistant/components/tautulli/sensor.py    | 18 ++++++++++--------
 3 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/homeassistant/components/tautulli/__init__.py b/homeassistant/components/tautulli/__init__.py
index b7fcf48cfdb..7d3efa4f283 100644
--- a/homeassistant/components/tautulli/__init__.py
+++ b/homeassistant/components/tautulli/__init__.py
@@ -16,9 +16,10 @@ from .const import DEFAULT_NAME, DOMAIN
 from .coordinator import TautulliDataUpdateCoordinator
 
 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."""
     host_configuration = PyTautulliHostConfiguration(
         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,
         session=async_get_clientsession(hass, entry.data[CONF_VERIFY_SSL]),
     )
-    coordinator = TautulliDataUpdateCoordinator(hass, host_configuration, api_client)
-    await coordinator.async_config_entry_first_refresh()
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+    entry.runtime_data = TautulliDataUpdateCoordinator(
+        hass, host_configuration, api_client
+    )
+    await entry.runtime_data.async_config_entry_first_refresh()
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     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."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
 class TautulliEntity(CoordinatorEntity[TautulliDataUpdateCoordinator]):
diff --git a/homeassistant/components/tautulli/coordinator.py b/homeassistant/components/tautulli/coordinator.py
index be7dfce4e3a..f392ab8df03 100644
--- a/homeassistant/components/tautulli/coordinator.py
+++ b/homeassistant/components/tautulli/coordinator.py
@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 from datetime import timedelta
+from typing import TYPE_CHECKING
 
 from pytautulli import (
     PyTautulli,
@@ -17,18 +18,20 @@ from pytautulli.exceptions import (
 )
 from pytautulli.models.host_configuration import PyTautulliHostConfiguration
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import DOMAIN, LOGGER
 
+if TYPE_CHECKING:
+    from . import TautulliConfigEntry
+
 
 class TautulliDataUpdateCoordinator(DataUpdateCoordinator[None]):
     """Data update coordinator for the Tautulli integration."""
 
-    config_entry: ConfigEntry
+    config_entry: TautulliConfigEntry
 
     def __init__(
         self,
diff --git a/homeassistant/components/tautulli/sensor.py b/homeassistant/components/tautulli/sensor.py
index f0d274bbe12..26b7c602de8 100644
--- a/homeassistant/components/tautulli/sensor.py
+++ b/homeassistant/components/tautulli/sensor.py
@@ -19,14 +19,14 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     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.core import HomeAssistant
 from homeassistant.helpers.entity import EntityDescription
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, StateType
 
-from . import TautulliEntity
+from . import TautulliConfigEntry, TautulliEntity
 from .const import ATTR_TOP_USER, DOMAIN
 from .coordinator import TautulliDataUpdateCoordinator
 
@@ -210,26 +210,28 @@ async def async_setup_platform(
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: TautulliConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Tautulli sensor."""
-    coordinator: TautulliDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
     entities: list[TautulliSensor | TautulliSessionSensor] = [
         TautulliSensor(
-            coordinator,
+            data,
             description,
         )
         for description in SENSOR_TYPES
     ]
-    if coordinator.users:
+    if data.users:
         entities.extend(
             TautulliSessionSensor(
-                coordinator,
+                data,
                 description,
                 user,
             )
             for description in SESSION_SENSOR_TYPES
-            for user in coordinator.users
+            for user in data.users
             if user.username != "Local"
         )
     async_add_entities(entities)