diff --git a/homeassistant/components/twentemilieu/__init__.py b/homeassistant/components/twentemilieu/__init__.py index d9881b0b2c8..b64a3ec2a1d 100644 --- a/homeassistant/components/twentemilieu/__init__.py +++ b/homeassistant/components/twentemilieu/__init__.py @@ -23,6 +23,9 @@ SERVICE_SCHEMA = vol.Schema({vol.Optional(CONF_ID): cv.string}) PLATFORMS = [Platform.CALENDAR, Platform.SENSOR] +TwenteMilieuDataUpdateCoordinator = DataUpdateCoordinator[dict[WasteType, list[date]]] +TwenteMilieuConfigEntry = ConfigEntry[TwenteMilieuDataUpdateCoordinator] + async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up Twente Milieu from a config entry.""" @@ -34,14 +37,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: session=session, ) - coordinator: DataUpdateCoordinator[dict[WasteType, list[date]]] = ( - DataUpdateCoordinator( - hass, - LOGGER, - name=DOMAIN, - update_interval=SCAN_INTERVAL, - update_method=twentemilieu.update, - ) + coordinator: TwenteMilieuDataUpdateCoordinator = DataUpdateCoordinator( + hass, + LOGGER, + name=DOMAIN, + update_interval=SCAN_INTERVAL, + update_method=twentemilieu.update, ) await coordinator.async_config_entry_first_refresh() @@ -51,7 +52,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: entry, unique_id=str(entry.data[CONF_ID]) ) - hass.data.setdefault(DOMAIN, {})[entry.data[CONF_ID]] = coordinator + entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True @@ -59,7 +60,4 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload Twente Milieu config entry.""" - unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - if unload_ok: - del hass.data[DOMAIN][entry.data[CONF_ID]] - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/twentemilieu/calendar.py b/homeassistant/components/twentemilieu/calendar.py index 8bd008e3eb3..8e7452823b7 100644 --- a/homeassistant/components/twentemilieu/calendar.py +++ b/homeassistant/components/twentemilieu/calendar.py @@ -2,30 +2,26 @@ from __future__ import annotations -from datetime import date, datetime, timedelta - -from twentemilieu import WasteType +from datetime import datetime, timedelta from homeassistant.components.calendar import CalendarEntity, CalendarEvent -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_ID from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator import homeassistant.util.dt as dt_util -from .const import DOMAIN, WASTE_TYPE_TO_DESCRIPTION +from . import TwenteMilieuConfigEntry +from .const import WASTE_TYPE_TO_DESCRIPTION from .entity import TwenteMilieuEntity async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: TwenteMilieuConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Twente Milieu calendar based on a config entry.""" - coordinator = hass.data[DOMAIN][entry.data[CONF_ID]] - async_add_entities([TwenteMilieuCalendar(coordinator, entry)]) + async_add_entities([TwenteMilieuCalendar(entry)]) class TwenteMilieuCalendar(TwenteMilieuEntity, CalendarEntity): @@ -35,13 +31,9 @@ class TwenteMilieuCalendar(TwenteMilieuEntity, CalendarEntity): _attr_name = None _attr_translation_key = "calendar" - def __init__( - self, - coordinator: DataUpdateCoordinator[dict[WasteType, list[date]]], - entry: ConfigEntry, - ) -> None: + def __init__(self, entry: TwenteMilieuConfigEntry) -> None: """Initialize the Twente Milieu entity.""" - super().__init__(coordinator, entry) + super().__init__(entry) self._attr_unique_id = str(entry.data[CONF_ID]) self._event: CalendarEvent | None = None diff --git a/homeassistant/components/twentemilieu/diagnostics.py b/homeassistant/components/twentemilieu/diagnostics.py index ea68473ae3b..9de3f9bfaff 100644 --- a/homeassistant/components/twentemilieu/diagnostics.py +++ b/homeassistant/components/twentemilieu/diagnostics.py @@ -2,29 +2,19 @@ from __future__ import annotations -from datetime import date from typing import Any -from twentemilieu import WasteType - from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_ID from homeassistant.core import HomeAssistant -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator - -from .const import DOMAIN async def async_get_config_entry_diagnostics( hass: HomeAssistant, entry: ConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - coordinator: DataUpdateCoordinator[dict[WasteType, list[date]]] = hass.data[DOMAIN][ - entry.data[CONF_ID] - ] return { f"WasteType.{waste_type.name}": [ waste_date.isoformat() for waste_date in waste_dates ] - for waste_type, waste_dates in coordinator.data.items() + for waste_type, waste_dates in entry.runtime_data.data.items() } diff --git a/homeassistant/components/twentemilieu/entity.py b/homeassistant/components/twentemilieu/entity.py index 1e0fa651998..896a8e32de9 100644 --- a/homeassistant/components/twentemilieu/entity.py +++ b/homeassistant/components/twentemilieu/entity.py @@ -2,36 +2,24 @@ from __future__ import annotations -from datetime import date - -from twentemilieu import WasteType - from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_ID from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.entity import Entity -from homeassistant.helpers.update_coordinator import ( - CoordinatorEntity, - DataUpdateCoordinator, -) +from homeassistant.helpers.update_coordinator import CoordinatorEntity +from . import TwenteMilieuDataUpdateCoordinator from .const import DOMAIN -class TwenteMilieuEntity( - CoordinatorEntity[DataUpdateCoordinator[dict[WasteType, list[date]]]], Entity -): +class TwenteMilieuEntity(CoordinatorEntity[TwenteMilieuDataUpdateCoordinator], Entity): """Defines a Twente Milieu entity.""" _attr_has_entity_name = True - def __init__( - self, - coordinator: DataUpdateCoordinator[dict[WasteType, list[date]]], - entry: ConfigEntry, - ) -> None: + def __init__(self, entry: ConfigEntry) -> None: """Initialize the Twente Milieu entity.""" - super().__init__(coordinator=coordinator) + super().__init__(coordinator=entry.runtime_data) self._attr_device_info = DeviceInfo( configuration_url="https://www.twentemilieu.nl", entry_type=DeviceEntryType.SERVICE, diff --git a/homeassistant/components/twentemilieu/sensor.py b/homeassistant/components/twentemilieu/sensor.py index f799fa62314..2d2e3de0f0e 100644 --- a/homeassistant/components/twentemilieu/sensor.py +++ b/homeassistant/components/twentemilieu/sensor.py @@ -16,7 +16,6 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_ID from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from .const import DOMAIN from .entity import TwenteMilieuEntity @@ -69,9 +68,8 @@ async def async_setup_entry( async_add_entities: AddEntitiesCallback, ) -> None: """Set up Twente Milieu sensor based on a config entry.""" - coordinator = hass.data[DOMAIN][entry.data[CONF_ID]] async_add_entities( - TwenteMilieuSensor(coordinator, description, entry) for description in SENSORS + TwenteMilieuSensor(entry, description) for description in SENSORS ) @@ -82,12 +80,11 @@ class TwenteMilieuSensor(TwenteMilieuEntity, SensorEntity): def __init__( self, - coordinator: DataUpdateCoordinator[dict[WasteType, list[date]]], - description: TwenteMilieuSensorDescription, entry: ConfigEntry, + description: TwenteMilieuSensorDescription, ) -> None: """Initialize the Twente Milieu entity.""" - super().__init__(coordinator, entry) + super().__init__(entry) self.entity_description = description self._attr_unique_id = f"{DOMAIN}_{entry.data[CONF_ID]}_{description.key}" diff --git a/tests/components/twentemilieu/test_init.py b/tests/components/twentemilieu/test_init.py index 901252f050f..d4c519d6f66 100644 --- a/tests/components/twentemilieu/test_init.py +++ b/tests/components/twentemilieu/test_init.py @@ -4,7 +4,6 @@ from unittest.mock import MagicMock, patch import pytest -from homeassistant.components.twentemilieu.const import DOMAIN from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant @@ -26,7 +25,6 @@ async def test_load_unload_config_entry( await hass.config_entries.async_unload(mock_config_entry.entry_id) await hass.async_block_till_done() - assert not hass.data.get(DOMAIN) assert mock_config_entry.state is ConfigEntryState.NOT_LOADED