Use ConfigEntry runtime_data in TwenteMilieu (#116642)

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]]] = (
+    coordinator: TwenteMilieuDataUpdateCoordinator = DataUpdateCoordinator(
-        DataUpdateCoordinator(
+        hass,
-            hass,
+        LOGGER,
-            LOGGER,
+        name=DOMAIN,
-            name=DOMAIN,
+        update_interval=SCAN_INTERVAL,
-            update_interval=SCAN_INTERVAL,
+        update_method=twentemilieu.update,
-            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)
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        del hass.data[DOMAIN][entry.data[CONF_ID]]
-    return unload_ok

homeassistant/components/twentemilieu/calendar.py

@@ -2,30 +2,26 @@
 
 from __future__ import annotations
 
-from datetime import date, datetime, timedelta
+from datetime import datetime, timedelta
-
-from twentemilieu import WasteType
 
 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(entry)])
-    async_add_entities([TwenteMilieuCalendar(coordinator, entry)])
 
 
 class TwenteMilieuCalendar(TwenteMilieuEntity, CalendarEntity):
@@ -35,13 +31,9 @@ class TwenteMilieuCalendar(TwenteMilieuEntity, CalendarEntity):
     _attr_name = None
     _attr_translation_key = "calendar"
 
-    def __init__(
+    def __init__(self, entry: TwenteMilieuConfigEntry) -> None:
-        self,
-        coordinator: DataUpdateCoordinator[dict[WasteType, list[date]]],
-        entry: ConfigEntry,
-    ) -> 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
 

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()
     }

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 (
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
 
+from . import TwenteMilieuDataUpdateCoordinator
 from .const import DOMAIN
 
 
-class TwenteMilieuEntity(
+class TwenteMilieuEntity(CoordinatorEntity[TwenteMilieuDataUpdateCoordinator], Entity):
-    CoordinatorEntity[DataUpdateCoordinator[dict[WasteType, list[date]]]], Entity
-):
     """Defines a Twente Milieu entity."""
 
     _attr_has_entity_name = True
 
-    def __init__(
+    def __init__(self, entry: ConfigEntry) -> None:
-        self,
-        coordinator: DataUpdateCoordinator[dict[WasteType, list[date]]],
-        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,

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}"
 

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