Use runtime_data for aftership (#116618)

homeassistant/components/aftership/__init__.py

@@ -10,16 +10,14 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import DOMAIN
-
 PLATFORMS: list[Platform] = [Platform.SENSOR]
 
+AfterShipConfigEntry = ConfigEntry[AfterShip]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: AfterShipConfigEntry) -> bool:
     """Set up AfterShip from a config entry."""
 
-    hass.data.setdefault(DOMAIN, {})
-
     session = async_get_clientsession(hass)
     aftership = AfterShip(api_key=entry.data[CONF_API_KEY], session=session)
 
@@ -28,7 +26,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     except AfterShipException as err:
         raise ConfigEntryNotReady from err
 
-    hass.data[DOMAIN][entry.entry_id] = aftership
+    entry.runtime_data = aftership
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -37,7 +35,4 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 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):
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/aftership/sensor.py

@@ -8,7 +8,6 @@ from typing import Any, Final
 from pyaftership import AfterShip, AfterShipException
 
 from homeassistant.components.sensor import SensorEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, ServiceCall
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.dispatcher import (
@@ -18,6 +17,7 @@ from homeassistant.helpers.dispatcher import (
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import Throttle
 
+from . import AfterShipConfigEntry
 from .const import (
     ADD_TRACKING_SERVICE_SCHEMA,
     ATTR_TRACKINGS,
@@ -41,11 +41,11 @@ PLATFORM_SCHEMA: Final = cv.removed(DOMAIN, raise_if_present=False)
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AfterShipConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up AfterShip sensor entities based on a config entry."""
-    aftership: AfterShip = hass.data[DOMAIN][config_entry.entry_id]
+    aftership = config_entry.runtime_data
 
     async_add_entities([AfterShipSensor(aftership, config_entry.title)], True)