From 9557ea902e05057b7074c68965a0b08b25e2f3f7 Mon Sep 17 00:00:00 2001
From: Michael <35783820+mib1185@users.noreply.github.com>
Date: Tue, 7 May 2024 22:13:53 +0200
Subject: [PATCH] Store runtime data inside the config entry in Apple TV
 (#117032)

store runtime data inside the config entry
---
 homeassistant/components/apple_tv/__init__.py     | 15 ++++++---------
 homeassistant/components/apple_tv/media_player.py |  8 +++-----
 homeassistant/components/apple_tv/remote.py       |  8 +++-----
 3 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/homeassistant/components/apple_tv/__init__.py b/homeassistant/components/apple_tv/__init__.py
index 5e3c1c37d4a..95bab5bc433 100644
--- a/homeassistant/components/apple_tv/__init__.py
+++ b/homeassistant/components/apple_tv/__init__.py
@@ -73,8 +73,10 @@ DEVICE_EXCEPTIONS = (
     exceptions.DeviceIdMissingError,
 )
 
+AppleTvConfigEntry = ConfigEntry["AppleTVManager"]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: AppleTvConfigEntry) -> bool:
     """Set up a config entry for Apple TV."""
     manager = AppleTVManager(hass, entry)
 
@@ -95,7 +97,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             )
             raise ConfigEntryNotReady(f"{address}: {ex}") from ex
 
-    hass.data.setdefault(DOMAIN, {})[entry.unique_id] = manager
+    entry.runtime_data = manager
 
     async def on_hass_stop(event: Event) -> None:
         """Stop push updates when hass stops."""
@@ -104,6 +106,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     entry.async_on_unload(
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, on_hass_stop)
     )
+    entry.async_on_unload(manager.disconnect)
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     await manager.init()
@@ -113,13 +116,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload an Apple TV config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-
-    if unload_ok:
-        manager = hass.data[DOMAIN].pop(entry.unique_id)
-        await manager.disconnect()
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
 class AppleTVEntity(Entity):
diff --git a/homeassistant/components/apple_tv/media_player.py b/homeassistant/components/apple_tv/media_player.py
index 3f64d10f9ac..9fb9dee46e1 100644
--- a/homeassistant/components/apple_tv/media_player.py
+++ b/homeassistant/components/apple_tv/media_player.py
@@ -37,15 +37,13 @@ from homeassistant.components.media_player import (
     RepeatMode,
     async_process_play_media_url,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.util.dt as dt_util
 
-from . import AppleTVEntity, AppleTVManager
+from . import AppleTvConfigEntry, AppleTVEntity, AppleTVManager
 from .browse_media import build_app_list
-from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -100,13 +98,13 @@ SUPPORT_FEATURE_MAPPING = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AppleTvConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Load Apple TV media player based on a config entry."""
     name: str = config_entry.data[CONF_NAME]
     assert config_entry.unique_id is not None
-    manager: AppleTVManager = hass.data[DOMAIN][config_entry.unique_id]
+    manager = config_entry.runtime_data
     async_add_entities([AppleTvMediaPlayer(name, config_entry.unique_id, manager)])
 
 
diff --git a/homeassistant/components/apple_tv/remote.py b/homeassistant/components/apple_tv/remote.py
index aed2c0ae3f0..8950a46388d 100644
--- a/homeassistant/components/apple_tv/remote.py
+++ b/homeassistant/components/apple_tv/remote.py
@@ -15,13 +15,11 @@ from homeassistant.components.remote import (
     DEFAULT_HOLD_SECS,
     RemoteEntity,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import AppleTVEntity, AppleTVManager
-from .const import DOMAIN
+from . import AppleTvConfigEntry, AppleTVEntity
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -38,14 +36,14 @@ COMMAND_TO_ATTRIBUTE = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AppleTvConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Load Apple TV remote based on a config entry."""
     name: str = config_entry.data[CONF_NAME]
     # apple_tv config entries always have a unique id
     assert config_entry.unique_id is not None
-    manager: AppleTVManager = hass.data[DOMAIN][config_entry.unique_id]
+    manager = config_entry.runtime_data
     async_add_entities([AppleTVRemote(name, config_entry.unique_id, manager)])