Migrate nzbget to use runtime_data (#166947)

homeassistant/components/nzbget/__init__.py

@@ -1,13 +1,12 @@
 """The NZBGet integration."""
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.typing import ConfigType
 
-from .const import DATA_COORDINATOR, DATA_UNDO_UPDATE_LISTENER, DOMAIN
-from .coordinator import NZBGetDataUpdateCoordinator
+from .const import DOMAIN
+from .coordinator import NZBGetConfigEntry, NZBGetDataUpdateCoordinator
 from .services import async_setup_services
 
 CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
@@ -22,37 +21,26 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: NZBGetConfigEntry) -> bool:
     """Set up NZBGet from a config entry."""
-    hass.data.setdefault(DOMAIN, {})
-
     coordinator = NZBGetDataUpdateCoordinator(hass, entry)
 
     await coordinator.async_config_entry_first_refresh()
 
-    undo_listener = entry.add_update_listener(_async_update_listener)
+    entry.runtime_data = coordinator
 
-    hass.data[DOMAIN][entry.entry_id] = {
-        DATA_COORDINATOR: coordinator,
-        DATA_UNDO_UPDATE_LISTENER: undo_listener,
-    }
+    entry.async_on_unload(entry.add_update_listener(_async_update_listener))
 
     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: NZBGetConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-
-    if unload_ok:
-        hass.data[DOMAIN][entry.entry_id][DATA_UNDO_UPDATE_LISTENER]()
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
-async def _async_update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def _async_update_listener(hass: HomeAssistant, entry: NZBGetConfigEntry) -> None:
     """Handle options update."""
     await hass.config_entries.async_reload(entry.entry_id)

homeassistant/components/nzbget/const.py

@@ -5,10 +5,6 @@ DOMAIN = "nzbget"
 # Attributes
 ATTR_SPEED = "speed"
 
-# Data
-DATA_COORDINATOR = "coordinator"
-DATA_UNDO_UPDATE_LISTENER = "undo_update_listener"
-
 # Defaults
 DEFAULT_NAME = "NZBGet"
 DEFAULT_PORT = 6789

homeassistant/components/nzbget/coordinator.py

@@ -23,15 +23,18 @@ from .const import DOMAIN
 _LOGGER = logging.getLogger(__name__)
 
 
+type NZBGetConfigEntry = ConfigEntry[NZBGetDataUpdateCoordinator]
+
+
 class NZBGetDataUpdateCoordinator(DataUpdateCoordinator):
     """Class to manage fetching NZBGet data."""
 
-    config_entry: ConfigEntry
+    config_entry: NZBGetConfigEntry
 
     def __init__(
         self,
         hass: HomeAssistant,
-        config_entry: ConfigEntry,
+        config_entry: NZBGetConfigEntry,
     ) -> None:
         """Initialize global NZBGet data updater."""
         self.nzbget = NZBGetAPI(

homeassistant/components/nzbget/sensor.py

@@ -10,15 +10,13 @@ from homeassistant.components.sensor import (
     SensorEntity,
     SensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME, UnitOfDataRate, UnitOfInformation
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.util.dt import utcnow
 
-from .const import DATA_COORDINATOR, DOMAIN
-from .coordinator import NZBGetDataUpdateCoordinator
+from .coordinator import NZBGetConfigEntry, NZBGetDataUpdateCoordinator
 from .entity import NZBGetEntity
 
 _LOGGER = logging.getLogger(__name__)
@@ -92,13 +90,11 @@ SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: NZBGetConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up NZBGet sensor based on a config entry."""
-    coordinator: NZBGetDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][
-        DATA_COORDINATOR
-    ]
+    coordinator = entry.runtime_data
     entities = [
         NZBGetSensor(coordinator, entry.entry_id, entry.data[CONF_NAME], description)
         for description in SENSOR_TYPES

homeassistant/components/nzbget/services.py

@@ -8,7 +8,6 @@ from homeassistant.helpers import config_validation as cv
 
 from .const import (
     ATTR_SPEED,
-    DATA_COORDINATOR,
     DEFAULT_SPEED_LIMIT,
     DOMAIN,
     SERVICE_PAUSE,
@@ -30,7 +29,7 @@ def _get_coordinator(call: ServiceCall) -> NZBGetDataUpdateCoordinator:
             translation_domain=DOMAIN,
             translation_key="invalid_config_entry",
         )
-    return call.hass.data[DOMAIN][entries[0].entry_id][DATA_COORDINATOR]
+    return entries[0].runtime_data
 
 
 def pause(call: ServiceCall) -> None:

homeassistant/components/nzbget/switch.py

@@ -5,25 +5,21 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.switch import SwitchEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import DATA_COORDINATOR, DOMAIN
-from .coordinator import NZBGetDataUpdateCoordinator
+from .coordinator import NZBGetConfigEntry, NZBGetDataUpdateCoordinator
 from .entity import NZBGetEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: NZBGetConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
-    """Set up NZBGet sensor based on a config entry."""
-    coordinator: NZBGetDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][
-        DATA_COORDINATOR
-    ]
+    """Set up NZBGet switch based on a config entry."""
+    coordinator = entry.runtime_data
 
     switches = [
         NZBGetDownloadSwitch(