Store runtime data inside the config entry in AsusWrt (#116889)

homeassistant/components/asuswrt/__init__.py

@@ -4,13 +4,14 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP, Platform
 from homeassistant.core import Event, HomeAssistant
 
-from .const import DATA_ASUSWRT, DOMAIN
 from .router import AsusWrtRouter
 
 PLATFORMS = [Platform.DEVICE_TRACKER, Platform.SENSOR]
 
+AsusWrtConfigEntry = ConfigEntry[AsusWrtRouter]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: AsusWrtConfigEntry) -> bool:
     """Set up AsusWrt platform."""
 
     router = AsusWrtRouter(hass, entry)
@@ -26,26 +27,25 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, async_close_connection)
     )
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {DATA_ASUSWRT: router}
+    entry.runtime_data = router
 
     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: AsusWrtConfigEntry) -> bool:
     """Unload a config entry."""
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        router = hass.data[DOMAIN][entry.entry_id][DATA_ASUSWRT]
+        router = entry.runtime_data
         await router.close()
-        hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
 
 
-async def update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def update_listener(hass: HomeAssistant, entry: AsusWrtConfigEntry) -> None:
     """Update when config_entry options update."""
-    router = hass.data[DOMAIN][entry.entry_id][DATA_ASUSWRT]
+    router = entry.runtime_data
 
     if router.update_options(entry.options):
         await hass.config_entries.async_reload(entry.entry_id)

homeassistant/components/asuswrt/const.py

@@ -8,8 +8,6 @@ CONF_REQUIRE_IP = "require_ip"
 CONF_SSH_KEY = "ssh_key"
 CONF_TRACK_UNKNOWN = "track_unknown"
 
-DATA_ASUSWRT = DOMAIN
-
 DEFAULT_DNSMASQ = "/var/lib/misc"
 DEFAULT_INTERFACE = "eth0"
 DEFAULT_TRACK_UNKNOWN = False

homeassistant/components/asuswrt/device_tracker.py

@@ -3,12 +3,11 @@
 from __future__ import annotations
 
 from homeassistant.components.device_tracker import ScannerEntity, SourceType
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DATA_ASUSWRT, DOMAIN
+from . import AsusWrtConfigEntry
 from .router import AsusWrtDevInfo, AsusWrtRouter
 
 ATTR_LAST_TIME_REACHABLE = "last_time_reachable"
@@ -17,10 +16,12 @@ DEFAULT_DEVICE_NAME = "Unknown device"
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: AsusWrtConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up device tracker for AsusWrt component."""
-    router = hass.data[DOMAIN][entry.entry_id][DATA_ASUSWRT]
+    router = entry.runtime_data
     tracked: set = set()
 
     @callback

homeassistant/components/asuswrt/diagnostics.py

@@ -7,7 +7,6 @@ from typing import Any
 import attr
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_CONNECTIONS,
     ATTR_IDENTIFIERS,
@@ -18,20 +17,19 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 
-from .const import DATA_ASUSWRT, DOMAIN
-from .router import AsusWrtRouter
+from . import AsusWrtConfigEntry
 
 TO_REDACT = {CONF_PASSWORD, CONF_UNIQUE_ID, CONF_USERNAME}
 TO_REDACT_DEV = {ATTR_CONNECTIONS, ATTR_IDENTIFIERS}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: AsusWrtConfigEntry
 ) -> dict[str, dict[str, Any]]:
     """Return diagnostics for a config entry."""
     data = {"entry": async_redact_data(entry.as_dict(), TO_REDACT)}
 
-    router: AsusWrtRouter = hass.data[DOMAIN][entry.entry_id][DATA_ASUSWRT]
+    router = entry.runtime_data
 
     # Gather information how this AsusWrt device is represented in Home Assistant
     device_registry = dr.async_get(hass)

homeassistant/components/asuswrt/router.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 from collections.abc import Callable
 from datetime import datetime, timedelta
 import logging
+from types import MappingProxyType
 from typing import Any
 
 from pyasuswrt import AsusWrtError
@@ -362,7 +363,7 @@ class AsusWrtRouter:
         """Add a function to call when router is closed."""
         self._on_close.append(func)
 
-    def update_options(self, new_options: dict[str, Any]) -> bool:
+    def update_options(self, new_options: MappingProxyType[str, Any]) -> bool:
         """Update router options."""
         req_reload = False
         for name, new_opt in new_options.items():

homeassistant/components/asuswrt/sensor.py

@@ -10,7 +10,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     EntityCategory,
     UnitOfDataRate,
@@ -25,9 +24,8 @@ from homeassistant.helpers.update_coordinator import (
 )
 from homeassistant.util import slugify
 
+from . import AsusWrtConfigEntry
 from .const import (
-    DATA_ASUSWRT,
-    DOMAIN,
     KEY_COORDINATOR,
     KEY_SENSORS,
     SENSORS_BYTES,
@@ -173,10 +171,12 @@ CONNECTION_SENSORS: tuple[AsusWrtSensorEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: AsusWrtConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the sensors."""
-    router: AsusWrtRouter = hass.data[DOMAIN][entry.entry_id][DATA_ASUSWRT]
+    router = entry.runtime_data
     entities = []
 
     for sensor_data in router.sensors_coordinator.values():