Store runtime data inside the config entry in Android TV Remote (#116824)

homeassistant/components/androidtv_remote/__init__.py

@@ -17,15 +17,18 @@ from homeassistant.const import CONF_HOST, CONF_NAME, EVENT_HOMEASSISTANT_STOP,
 from homeassistant.core import Event, HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 
-from .const import DOMAIN
 from .helpers import create_api, get_enable_ime
 
 _LOGGER = logging.getLogger(__name__)
 
 PLATFORMS: list[Platform] = [Platform.MEDIA_PLAYER, Platform.REMOTE]
 
+AndroidTVRemoteConfigEntry = ConfigEntry[AndroidTVRemote]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(
+    hass: HomeAssistant, entry: AndroidTVRemoteConfigEntry
+) -> bool:
     """Set up Android TV Remote from a config entry."""
     api = create_api(hass, entry.data[CONF_HOST], get_enable_ime(entry))
 
@@ -64,7 +67,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     # update the config entry data and reload the config entry.
     api.keep_reconnecting(reauth_needed)
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = api
+    entry.runtime_data = api
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -77,17 +80,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, on_hass_stop)
     )
     entry.async_on_unload(entry.add_update_listener(update_listener))
+    entry.async_on_unload(api.disconnect)
 
     return True
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """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)
-        api: AndroidTVRemote = hass.data[DOMAIN].pop(entry.entry_id)
-        api.disconnect()
-
-    return unload_ok
 
 
 async def update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:

homeassistant/components/androidtv_remote/diagnostics.py

@@ -4,23 +4,20 @@ from __future__ import annotations
 
 from typing import Any
 
-from androidtvremote2 import AndroidTVRemote
-
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_MAC
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
+from . import AndroidTVRemoteConfigEntry
 
 TO_REDACT = {CONF_HOST, CONF_MAC}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: AndroidTVRemoteConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    api: AndroidTVRemote = hass.data[DOMAIN].pop(entry.entry_id)
+    api = entry.runtime_data
     return async_redact_data(
         {
             "api_device_info": api.device_info,

homeassistant/components/androidtv_remote/media_player.py

@@ -14,12 +14,11 @@ from homeassistant.components.media_player import (
     MediaPlayerState,
     MediaType,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import AndroidTVRemoteConfigEntry
 from .entity import AndroidTVRemoteBaseEntity
 
 PARALLEL_UPDATES = 0
@@ -27,11 +26,11 @@ PARALLEL_UPDATES = 0
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AndroidTVRemoteConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Android TV media player entity based on a config entry."""
-    api: AndroidTVRemote = hass.data[DOMAIN][config_entry.entry_id]
+    api = config_entry.runtime_data
     async_add_entities([AndroidTVRemoteMediaPlayerEntity(api, config_entry)])
 
 
@@ -53,7 +52,9 @@ class AndroidTVRemoteMediaPlayerEntity(AndroidTVRemoteBaseEntity, MediaPlayerEnt
         | MediaPlayerEntityFeature.PLAY_MEDIA
     )
 
-    def __init__(self, api: AndroidTVRemote, config_entry: ConfigEntry) -> None:
+    def __init__(
+        self, api: AndroidTVRemote, config_entry: AndroidTVRemoteConfigEntry
+    ) -> None:
         """Initialize the entity."""
         super().__init__(api, config_entry)
 

homeassistant/components/androidtv_remote/remote.py

@@ -6,8 +6,6 @@ import asyncio
 from collections.abc import Iterable
 from typing import Any
 
-from androidtvremote2 import AndroidTVRemote
-
 from homeassistant.components.remote import (
     ATTR_ACTIVITY,
     ATTR_DELAY_SECS,
@@ -19,11 +17,10 @@ from homeassistant.components.remote import (
     RemoteEntity,
     RemoteEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import AndroidTVRemoteConfigEntry
 from .entity import AndroidTVRemoteBaseEntity
 
 PARALLEL_UPDATES = 0
@@ -31,11 +28,11 @@ PARALLEL_UPDATES = 0
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AndroidTVRemoteConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Android TV remote entity based on a config entry."""
-    api: AndroidTVRemote = hass.data[DOMAIN][config_entry.entry_id]
+    api = config_entry.runtime_data
     async_add_entities([AndroidTVRemoteEntity(api, config_entry)])