diff --git a/homeassistant/components/dwd_weather_warnings/__init__.py b/homeassistant/components/dwd_weather_warnings/__init__.py index 209c77f60b5..f71b81d862b 100644 --- a/homeassistant/components/dwd_weather_warnings/__init__.py +++ b/homeassistant/components/dwd_weather_warnings/__init__.py @@ -2,27 +2,27 @@ from __future__ import annotations -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from .const import DOMAIN, PLATFORMS -from .coordinator import DwdWeatherWarningsCoordinator +from .const import PLATFORMS +from .coordinator import DwdWeatherWarningsConfigEntry, DwdWeatherWarningsCoordinator -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry( + hass: HomeAssistant, entry: DwdWeatherWarningsConfigEntry +) -> bool: """Set up a config entry.""" coordinator = DwdWeatherWarningsCoordinator(hass) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + entry.runtime_data = coordinator 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: DwdWeatherWarningsConfigEntry +) -> 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/dwd_weather_warnings/coordinator.py b/homeassistant/components/dwd_weather_warnings/coordinator.py index 7600a04f2bb..7f0afe352db 100644 --- a/homeassistant/components/dwd_weather_warnings/coordinator.py +++ b/homeassistant/components/dwd_weather_warnings/coordinator.py @@ -19,11 +19,13 @@ from .const import ( from .exceptions import EntityNotFoundError from .util import get_position_data +DwdWeatherWarningsConfigEntry = ConfigEntry["DwdWeatherWarningsCoordinator"] + class DwdWeatherWarningsCoordinator(DataUpdateCoordinator[None]): """Custom coordinator for the dwd_weather_warnings integration.""" - config_entry: ConfigEntry + config_entry: DwdWeatherWarningsConfigEntry api: DwdWeatherWarningsAPI def __init__(self, hass: HomeAssistant) -> None: diff --git a/homeassistant/components/dwd_weather_warnings/sensor.py b/homeassistant/components/dwd_weather_warnings/sensor.py index d62c0f4f192..cef665ffb10 100644 --- a/homeassistant/components/dwd_weather_warnings/sensor.py +++ b/homeassistant/components/dwd_weather_warnings/sensor.py @@ -14,7 +14,6 @@ from __future__ import annotations from typing import Any from homeassistant.components.sensor import SensorEntity, SensorEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback @@ -40,7 +39,7 @@ from .const import ( DEFAULT_NAME, DOMAIN, ) -from .coordinator import DwdWeatherWarningsCoordinator +from .coordinator import DwdWeatherWarningsConfigEntry, DwdWeatherWarningsCoordinator SENSOR_TYPES: tuple[SensorEntityDescription, ...] = ( SensorEntityDescription( @@ -55,10 +54,12 @@ SENSOR_TYPES: tuple[SensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: DwdWeatherWarningsConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up entities from config entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( [ @@ -80,7 +81,7 @@ class DwdWeatherWarningsSensor( def __init__( self, coordinator: DwdWeatherWarningsCoordinator, - entry: ConfigEntry, + entry: DwdWeatherWarningsConfigEntry, description: SensorEntityDescription, ) -> None: """Initialize a DWD-Weather-Warnings sensor.""" diff --git a/tests/components/dwd_weather_warnings/test_init.py b/tests/components/dwd_weather_warnings/test_init.py index 360efc390db..e5b82d0c453 100644 --- a/tests/components/dwd_weather_warnings/test_init.py +++ b/tests/components/dwd_weather_warnings/test_init.py @@ -6,6 +6,9 @@ from homeassistant.components.dwd_weather_warnings.const import ( CONF_REGION_DEVICE_TRACKER, DOMAIN, ) +from homeassistant.components.dwd_weather_warnings.coordinator import ( + DwdWeatherWarningsCoordinator, +) from homeassistant.config_entries import ConfigEntryState from homeassistant.const import ATTR_LATITUDE, ATTR_LONGITUDE, STATE_HOME from homeassistant.core import HomeAssistant @@ -25,13 +28,12 @@ async def test_load_unload_entry( entry = await init_integration(hass, mock_identifier_entry) assert entry.state is ConfigEntryState.LOADED - assert entry.entry_id in hass.data[DOMAIN] + assert isinstance(entry.runtime_data, DwdWeatherWarningsCoordinator) assert await hass.config_entries.async_unload(entry.entry_id) await hass.async_block_till_done() assert entry.state is ConfigEntryState.NOT_LOADED - assert entry.entry_id not in hass.data[DOMAIN] async def test_load_invalid_registry_entry( @@ -97,4 +99,4 @@ async def test_load_valid_device_tracker( await hass.async_block_till_done() assert mock_tracker_entry.state is ConfigEntryState.LOADED - assert mock_tracker_entry.entry_id in hass.data[DOMAIN] + assert isinstance(mock_tracker_entry.runtime_data, DwdWeatherWarningsCoordinator)