Migrate notion to use runtime_data (#166936)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

homeassistant/components/notion/__init__.py

@@ -9,7 +9,6 @@ from uuid import UUID
 from aionotion.errors import InvalidCredentialsError, NotionError
 from aionotion.listener.models import ListenerKind
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
@@ -18,7 +17,6 @@ from homeassistant.helpers import entity_registry as er
 from .const import (
     CONF_REFRESH_TOKEN,
     CONF_USER_UUID,
-    DOMAIN,
     LOGGER,
     SENSOR_BATTERY,
     SENSOR_DOOR,
@@ -31,7 +29,7 @@ from .const import (
     SENSOR_TEMPERATURE,
     SENSOR_WINDOW_HINGED,
 )
-from .coordinator import NotionDataUpdateCoordinator
+from .coordinator import NotionConfigEntry, NotionDataUpdateCoordinator
 from .util import async_get_client_with_credentials, async_get_client_with_refresh_token
 
 PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR]
@@ -67,7 +65,7 @@ def is_uuid(value: str) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: NotionConfigEntry) -> bool:
     """Set up Notion as a config entry."""
     entry_updates: dict[str, Any] = {"data": {**entry.data}}
 
@@ -119,8 +117,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     coordinator = NotionDataUpdateCoordinator(hass, entry=entry, client=client)
 
     await coordinator.async_config_entry_first_refresh()
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
 
     @callback
     def async_migrate_entity_entry(entry: er.RegistryEntry) -> dict[str, Any] | None:
@@ -157,10 +154,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: NotionConfigEntry) -> bool:
     """Unload a Notion config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/notion/binary_sensor.py

@@ -12,13 +12,11 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
 from .const import (
-    DOMAIN,
     LOGGER,
     SENSOR_BATTERY,
     SENSOR_DOOR,
@@ -30,7 +28,7 @@ from .const import (
     SENSOR_SMOKE_CO,
     SENSOR_WINDOW_HINGED,
 )
-from .coordinator import NotionDataUpdateCoordinator
+from .coordinator import NotionConfigEntry
 from .entity import NotionEntity, NotionEntityDescription
 
 
@@ -108,11 +106,11 @@ BINARY_SENSOR_DESCRIPTIONS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: NotionConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Notion sensors based on a config entry."""
-    coordinator: NotionDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     async_add_entities(
         [

homeassistant/components/notion/coordinator.py

@@ -28,10 +28,12 @@ DATA_USER_PREFERENCES = "user_preferences"
 
 DEFAULT_SCAN_INTERVAL = timedelta(minutes=1)
 
+type NotionConfigEntry = ConfigEntry[NotionDataUpdateCoordinator]
+
 
 @callback
 def _async_register_new_bridge(
-    hass: HomeAssistant, entry: ConfigEntry, bridge: Bridge
+    hass: HomeAssistant, entry: NotionConfigEntry, bridge: Bridge
 ) -> None:
     """Register a new bridge."""
     if name := bridge.name:
@@ -55,7 +57,7 @@ class NotionData:
     """Define a manager class for Notion data."""
 
     hass: HomeAssistant
-    entry: ConfigEntry
+    entry: NotionConfigEntry
 
     # Define a dict of bridges, indexed by bridge ID (an integer):
     bridges: dict[int, Bridge] = field(default_factory=dict)
@@ -104,13 +106,13 @@ class NotionData:
 class NotionDataUpdateCoordinator(DataUpdateCoordinator[NotionData]):
     """Define a Notion data coordinator."""
 
-    config_entry: ConfigEntry
+    config_entry: NotionConfigEntry
 
     def __init__(
         self,
         hass: HomeAssistant,
         *,
-        entry: ConfigEntry,
+        entry: NotionConfigEntry,
         client: Client,
     ) -> None:
         """Initialize."""

homeassistant/components/notion/diagnostics.py

@@ -5,12 +5,11 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_EMAIL, CONF_UNIQUE_ID, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 
-from .const import CONF_REFRESH_TOKEN, CONF_USER_UUID, DOMAIN
-from .coordinator import NotionDataUpdateCoordinator
+from .const import CONF_REFRESH_TOKEN, CONF_USER_UUID
+from .coordinator import NotionConfigEntry
 
 CONF_DEVICE_KEY = "device_key"
 CONF_HARDWARE_ID = "hardware_id"
@@ -34,10 +33,10 @@ TO_REDACT = {
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: NotionConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: NotionDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     return async_redact_data(
         {

homeassistant/components/notion/sensor.py

@@ -10,13 +10,12 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import DOMAIN, SENSOR_MOLD, SENSOR_TEMPERATURE
-from .coordinator import NotionDataUpdateCoordinator
+from .const import SENSOR_MOLD, SENSOR_TEMPERATURE
+from .coordinator import NotionConfigEntry
 from .entity import NotionEntity, NotionEntityDescription
 
 
@@ -43,11 +42,11 @@ SENSOR_DESCRIPTIONS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: NotionConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Notion sensors based on a config entry."""
-    coordinator: NotionDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     async_add_entities(
         [

tests/components/notion/conftest.py

@@ -10,7 +10,11 @@ from aionotion.sensor.models import Sensor
 from aionotion.user.models import UserPreferences
 import pytest
 
-from homeassistant.components.notion import CONF_REFRESH_TOKEN, CONF_USER_UUID, DOMAIN
+from homeassistant.components.notion.const import (
+    CONF_REFRESH_TOKEN,
+    CONF_USER_UUID,
+    DOMAIN,
+)
 from homeassistant.const import CONF_USERNAME
 from homeassistant.core import HomeAssistant
 

tests/components/notion/test_config_flow.py

@@ -5,7 +5,11 @@ from unittest.mock import AsyncMock, patch
 from aionotion.errors import InvalidCredentialsError, NotionError
 import pytest
 
-from homeassistant.components.notion import CONF_REFRESH_TOKEN, CONF_USER_UUID, DOMAIN
+from homeassistant.components.notion.const import (
+    CONF_REFRESH_TOKEN,
+    CONF_USER_UUID,
+    DOMAIN,
+)
 from homeassistant.config_entries import SOURCE_USER
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant

tests/components/notion/test_diagnostics.py

@@ -1,7 +1,7 @@
 """Test Notion diagnostics."""
 
 from homeassistant.components.diagnostics import REDACTED
-from homeassistant.components.notion import DOMAIN
+from homeassistant.components.notion.const import DOMAIN
 from homeassistant.core import HomeAssistant
 
 from tests.common import ANY