diff --git a/homeassistant/components/aftership/__init__.py b/homeassistant/components/aftership/__init__.py index b079079db08..10e4293bc51 100644 --- a/homeassistant/components/aftership/__init__.py +++ b/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) diff --git a/homeassistant/components/aftership/sensor.py b/homeassistant/components/aftership/sensor.py index c403c4a571d..c019634197d 100644 --- a/homeassistant/components/aftership/sensor.py +++ b/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)