Use runtime_data in sabnzbd (#131069)

homeassistant/components/sabnzbd/__init__.py

@@ -48,20 +48,24 @@ SERVICE_SPEED_SCHEMA = SERVICE_BASE_SCHEMA.extend(
     }
 )
 
+type SabnzbdConfigEntry = ConfigEntry[SabnzbdUpdateCoordinator]
+
 
 @callback
-def async_get_entry_id_for_service_call(hass: HomeAssistant, call: ServiceCall) -> str:
+def async_get_entry_for_service_call(
+    hass: HomeAssistant, call: ServiceCall
+) -> SabnzbdConfigEntry:
     """Get the entry ID related to a service call (by device ID)."""
     call_data_api_key = call.data[ATTR_API_KEY]
 
     for entry in hass.config_entries.async_entries(DOMAIN):
         if entry.data[ATTR_API_KEY] == call_data_api_key:
-            return entry.entry_id
+            return entry
 
     raise ValueError(f"No api for API key: {call_data_api_key}")
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: SabnzbdConfigEntry) -> bool:
     """Set up the SabNzbd Component."""
 
     sab_api = await get_client(hass, entry.data)
@@ -70,7 +74,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     coordinator = SabnzbdUpdateCoordinator(hass, entry, sab_api)
     await coordinator.async_config_entry_first_refresh()
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
 
     @callback
     def extract_api(
@@ -82,8 +86,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
         async def wrapper(call: ServiceCall) -> None:
             """Wrap the service function."""
-            entry_id = async_get_entry_id_for_service_call(hass, call)
+            config_entry = async_get_entry_for_service_call(hass, call)
-            coordinator: SabnzbdUpdateCoordinator = hass.data[DOMAIN][entry_id]
+            coordinator = config_entry.runtime_data
 
             try:
                 await func(call, coordinator)
@@ -155,11 +159,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: SabnzbdConfigEntry) -> bool:
     """Unload a Sabnzbd config entry."""
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
 
     loaded_entries = [
         entry

homeassistant/components/sabnzbd/button.py

@@ -7,13 +7,13 @@ from typing import Any
 from pysabnzbd import SabnzbdApiException
 
 from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import SabnzbdUpdateCoordinator
+from . import SabnzbdConfigEntry
 from .const import DOMAIN
+from .coordinator import SabnzbdUpdateCoordinator
 from .entity import SabnzbdEntity
 
 
@@ -40,11 +40,11 @@ BUTTON_DESCRIPTIONS: tuple[SabnzbdButtonEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: SabnzbdConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up buttons from a config entry."""
-    coordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     async_add_entities(
         SabnzbdButton(coordinator, description) for description in BUTTON_DESCRIPTIONS

homeassistant/components/sabnzbd/number.py

@@ -12,12 +12,12 @@ from homeassistant.components.number import (
     NumberEntityDescription,
     NumberMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from . import SabnzbdConfigEntry
 from .const import DOMAIN
 from .coordinator import SabnzbdUpdateCoordinator
 from .entity import SabnzbdEntity
@@ -48,11 +48,11 @@ NUMBER_DESCRIPTIONS: tuple[SabnzbdNumberEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: SabnzbdConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the SABnzbd number entity."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data
 
     async_add_entities(
         SabnzbdNumber(coordinator, description) for description in NUMBER_DESCRIPTIONS

homeassistant/components/sabnzbd/sensor.py

@@ -10,14 +10,12 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfDataRate, UnitOfInformation
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
-from .const import DOMAIN
+from . import SabnzbdConfigEntry
-from .coordinator import SabnzbdUpdateCoordinator
 from .entity import SabnzbdEntity
 
 
@@ -116,13 +114,11 @@ SENSOR_TYPES: tuple[SabnzbdSensorEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: SabnzbdConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up a Sabnzbd sensor entry."""
-
+    coordinator = config_entry.runtime_data
-    entry_id = config_entry.entry_id
-    coordinator: SabnzbdUpdateCoordinator = hass.data[DOMAIN][entry_id]
 
     async_add_entities([SabnzbdSensor(coordinator, sensor) for sensor in SENSOR_TYPES])