Migrate kodi to use runtime_data (#147191)

homeassistant/components/kodi/__init__.py

@@ -1,8 +1,10 @@
 """The kodi component."""
 
+from dataclasses import dataclass
 import logging
 
 from pykodi import CannotConnectError, InvalidAuthError, Kodi, get_kodi_connection
+from pykodi.kodi import KodiHTTPConnection, KodiWSConnection
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
@@ -17,13 +19,23 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import CONF_WS_PORT, DATA_CONNECTION, DATA_KODI, DOMAIN
+from .const import CONF_WS_PORT
 
 _LOGGER = logging.getLogger(__name__)
 PLATFORMS = [Platform.MEDIA_PLAYER]
 
+type KodiConfigEntry = ConfigEntry[KodiRuntimeData]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+@dataclass
+class KodiRuntimeData:
+    """Data class to hold Kodi runtime data."""
+
+    connection: KodiHTTPConnection | KodiWSConnection
+    kodi: Kodi
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: KodiConfigEntry) -> bool:
     """Set up Kodi from a config entry."""
     conn = get_kodi_connection(
         entry.data[CONF_HOST],
@@ -54,22 +66,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     entry.async_on_unload(hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _close))
 
-    hass.data.setdefault(DOMAIN, {})
+    entry.runtime_data = KodiRuntimeData(connection=conn, kodi=kodi)
-    hass.data[DOMAIN][entry.entry_id] = {
-        DATA_CONNECTION: conn,
-        DATA_KODI: kodi,
-    }
 
     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: KodiConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-    if unload_ok:
+        await entry.runtime_data.connection.close()
-        data = hass.data[DOMAIN].pop(entry.entry_id)
-        await data[DATA_CONNECTION].close()
 
     return unload_ok

homeassistant/components/kodi/const.py

@@ -4,9 +4,6 @@ DOMAIN = "kodi"
 
 CONF_WS_PORT = "ws_port"
 
-DATA_CONNECTION = "connection"
-DATA_KODI = "kodi"
-
 DEFAULT_PORT = 8080
 DEFAULT_SSL = False
 DEFAULT_TIMEOUT = 5

homeassistant/components/kodi/media_player.py

@@ -24,7 +24,7 @@ from homeassistant.components.media_player import (
     MediaType,
     async_process_play_media_url,
 )
-from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
+from homeassistant.config_entries import SOURCE_IMPORT
 from homeassistant.const import (
     ATTR_ENTITY_ID,
     CONF_DEVICE_ID,
@@ -55,6 +55,7 @@ from homeassistant.helpers.network import is_internal_request
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, VolDictType
 from homeassistant.util import dt as dt_util
 
+from . import KodiConfigEntry
 from .browse_media import (
     build_item_response,
     get_media_info,
@@ -63,8 +64,6 @@ from .browse_media import (
 )
 from .const import (
     CONF_WS_PORT,
-    DATA_CONNECTION,
-    DATA_KODI,
     DEFAULT_PORT,
     DEFAULT_SSL,
     DEFAULT_TIMEOUT,
@@ -208,7 +207,7 @@ async def async_setup_platform(
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: KodiConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the Kodi media player platform."""
@@ -220,14 +219,12 @@ async def async_setup_entry(
         SERVICE_CALL_METHOD, KODI_CALL_METHOD_SCHEMA, "async_call_method"
     )
 
-    data = hass.data[DOMAIN][config_entry.entry_id]
+    data = config_entry.runtime_data
-    connection = data[DATA_CONNECTION]
-    kodi = data[DATA_KODI]
     name = config_entry.data[CONF_NAME]
     if (uid := config_entry.unique_id) is None:
         uid = config_entry.entry_id
 
-    entity = KodiEntity(connection, kodi, name, uid)
+    entity = KodiEntity(data.connection, data.kodi, name, uid)
     async_add_entities([entity])
 
 

tests/components/kodi/test_device_trigger.py

@@ -4,7 +4,7 @@ import pytest
 
 from homeassistant.components import automation
 from homeassistant.components.device_automation import DeviceAutomationType
-from homeassistant.components.kodi import DOMAIN
+from homeassistant.components.kodi.const import DOMAIN
 from homeassistant.components.media_player import DOMAIN as MP_DOMAIN
 from homeassistant.core import HomeAssistant, ServiceCall
 from homeassistant.helpers import device_registry as dr, entity_registry as er