diff --git a/homeassistant/components/here_travel_time/__init__.py b/homeassistant/components/here_travel_time/__init__.py index 132b12de4ce..525da15bd74 100644 --- a/homeassistant/components/here_travel_time/__init__.py +++ b/homeassistant/components/here_travel_time/__init__.py @@ -2,7 +2,6 @@ from __future__ import annotations -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_API_KEY, CONF_MODE, Platform from homeassistant.core import HomeAssistant from homeassistant.helpers.start import async_at_started @@ -18,10 +17,10 @@ from .const import ( CONF_ORIGIN_LATITUDE, CONF_ORIGIN_LONGITUDE, CONF_ROUTE_MODE, - DOMAIN, TRAVEL_MODE_PUBLIC, ) from .coordinator import ( + HereConfigEntry, HERERoutingDataUpdateCoordinator, HERETransitDataUpdateCoordinator, ) @@ -30,7 +29,7 @@ from .model import HERETravelTimeConfig PLATFORMS = [Platform.SENSOR] -async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, config_entry: HereConfigEntry) -> bool: """Set up HERE Travel Time from a config entry.""" api_key = config_entry.data[CONF_API_KEY] @@ -57,7 +56,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b cls = HERERoutingDataUpdateCoordinator data_coordinator = cls(hass, config_entry, api_key, here_travel_time_config) - hass.data.setdefault(DOMAIN, {})[config_entry.entry_id] = data_coordinator + config_entry.runtime_data = data_coordinator async def _async_update_at_start(_: HomeAssistant) -> None: await data_coordinator.async_refresh() @@ -68,12 +67,8 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b return True -async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: +async def async_unload_entry( + hass: HomeAssistant, config_entry: HereConfigEntry +) -> bool: """Unload a config entry.""" - unload_ok = await hass.config_entries.async_unload_platforms( - config_entry, PLATFORMS - ) - if unload_ok: - hass.data[DOMAIN].pop(config_entry.entry_id) - - return unload_ok + return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS) diff --git a/homeassistant/components/here_travel_time/coordinator.py b/homeassistant/components/here_travel_time/coordinator.py index a3345e78e4e..aa36404c584 100644 --- a/homeassistant/components/here_travel_time/coordinator.py +++ b/homeassistant/components/here_travel_time/coordinator.py @@ -41,16 +41,20 @@ BACKOFF_MULTIPLIER = 1.1 _LOGGER = logging.getLogger(__name__) +type HereConfigEntry = ConfigEntry[ + HERETransitDataUpdateCoordinator | HERERoutingDataUpdateCoordinator +] + class HERERoutingDataUpdateCoordinator(DataUpdateCoordinator[HERETravelTimeData]): """here_routing DataUpdateCoordinator.""" - config_entry: ConfigEntry + config_entry: HereConfigEntry def __init__( self, hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: HereConfigEntry, api_key: str, config: HERETravelTimeConfig, ) -> None: @@ -173,12 +177,12 @@ class HERETransitDataUpdateCoordinator( ): """HERETravelTime DataUpdateCoordinator.""" - config_entry: ConfigEntry + config_entry: HereConfigEntry def __init__( self, hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: HereConfigEntry, api_key: str, config: HERETravelTimeConfig, ) -> None: diff --git a/homeassistant/components/here_travel_time/sensor.py b/homeassistant/components/here_travel_time/sensor.py index 0f0cbb7d3cb..bbaabb56d46 100644 --- a/homeassistant/components/here_travel_time/sensor.py +++ b/homeassistant/components/here_travel_time/sensor.py @@ -12,7 +12,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( ATTR_ATTRIBUTION, ATTR_LATITUDE, @@ -40,6 +39,7 @@ from .const import ( ICONS, ) from .coordinator import ( + HereConfigEntry, HERERoutingDataUpdateCoordinator, HERETransitDataUpdateCoordinator, ) @@ -77,14 +77,14 @@ def sensor_descriptions(travel_mode: str) -> tuple[SensorEntityDescription, ...] async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: HereConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Add HERE travel time entities from a config_entry.""" entry_id = config_entry.entry_id name = config_entry.data[CONF_NAME] - coordinator = hass.data[DOMAIN][entry_id] + coordinator = config_entry.runtime_data sensors: list[HERETravelTimeSensor] = [ HERETravelTimeSensor( @@ -164,7 +164,8 @@ class OriginSensor(HERETravelTimeSensor): self, unique_id_prefix: str, name: str, - coordinator: HERERoutingDataUpdateCoordinator, + coordinator: HERERoutingDataUpdateCoordinator + | HERETransitDataUpdateCoordinator, ) -> None: """Initialize the sensor.""" sensor_description = SensorEntityDescription( @@ -192,7 +193,8 @@ class DestinationSensor(HERETravelTimeSensor): self, unique_id_prefix: str, name: str, - coordinator: HERERoutingDataUpdateCoordinator, + coordinator: HERERoutingDataUpdateCoordinator + | HERETransitDataUpdateCoordinator, ) -> None: """Initialize the sensor.""" sensor_description = SensorEntityDescription( diff --git a/tests/components/here_travel_time/test_init.py b/tests/components/here_travel_time/test_init.py index 682d8c560bb..ff09c7e6ae9 100644 --- a/tests/components/here_travel_time/test_init.py +++ b/tests/components/here_travel_time/test_init.py @@ -50,4 +50,3 @@ async def test_unload_entry(hass: HomeAssistant, options) -> None: await hass.config_entries.async_setup(entry.entry_id) await hass.async_block_till_done() assert await hass.config_entries.async_unload(entry.entry_id) - assert not hass.data[DOMAIN]