Migrate keenetic_ndms2 to use runtime_data (#147194)

* Migrate keenetic_ndms2 to use runtime_data

* Adjust tests

homeassistant/components/keenetic_ndms2/__init__.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 import logging
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_SCAN_INTERVAL, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr, entity_registry as er
@@ -19,15 +18,14 @@ from .const import (
     DEFAULT_INTERFACE,
     DEFAULT_SCAN_INTERVAL,
     DOMAIN,
-    ROUTER,
 )
-from .router import KeeneticRouter
+from .router import KeeneticConfigEntry, KeeneticRouter
 
 PLATFORMS = [Platform.BINARY_SENSOR, Platform.DEVICE_TRACKER]
 _LOGGER = logging.getLogger(__name__)
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: KeeneticConfigEntry) -> bool:
     """Set up the component."""
     hass.data.setdefault(DOMAIN, {})
     async_add_defaults(hass, entry)
@@ -37,27 +35,24 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     entry.async_on_unload(entry.add_update_listener(update_listener))
 
-    hass.data[DOMAIN][entry.entry_id] = {
-        ROUTER: router,
-    }
+    entry.runtime_data = router
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_unload_entry(
+    hass: HomeAssistant, config_entry: KeeneticConfigEntry
+) -> bool:
     """Unload a config entry."""
     unload_ok = await hass.config_entries.async_unload_platforms(
         config_entry, PLATFORMS
     )
 
-    router: KeeneticRouter = hass.data[DOMAIN][config_entry.entry_id][ROUTER]
-
+    router = config_entry.runtime_data
     await router.async_teardown()
 
-    hass.data[DOMAIN].pop(config_entry.entry_id)
-
     new_tracked_interfaces: set[str] = set(config_entry.options[CONF_INTERFACES])
 
     if router.tracked_interfaces - new_tracked_interfaces:
@@ -92,12 +87,12 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
     return unload_ok
 
 
-async def update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def update_listener(hass: HomeAssistant, entry: KeeneticConfigEntry) -> None:
     """Handle options update."""
     await hass.config_entries.async_reload(entry.entry_id)
 
 
-def async_add_defaults(hass: HomeAssistant, entry: ConfigEntry):
+def async_add_defaults(hass: HomeAssistant, entry: KeeneticConfigEntry):
     """Populate default options."""
     host: str = entry.data[CONF_HOST]
     imported_options: dict = hass.data[DOMAIN].get(f"imported_options_{host}", {})

homeassistant/components/keenetic_ndms2/binary_sensor.py

@@ -4,24 +4,20 @@ from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from . import KeeneticRouter
-from .const import DOMAIN, ROUTER
+from .router import KeeneticConfigEntry, KeeneticRouter
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: KeeneticConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up device tracker for Keenetic NDMS2 component."""
-    router: KeeneticRouter = hass.data[DOMAIN][config_entry.entry_id][ROUTER]
-
-    async_add_entities([RouterOnlineBinarySensor(router)])
+    async_add_entities([RouterOnlineBinarySensor(config_entry.runtime_data)])
 
 
 class RouterOnlineBinarySensor(BinarySensorEntity):

homeassistant/components/keenetic_ndms2/config_flow.py

@@ -8,12 +8,7 @@ from urllib.parse import urlparse
 from ndms2_client import Client, ConnectionException, InterfaceInfo, TelnetConnection
 import voluptuous as vol
 
-from homeassistant.config_entries import (
-    ConfigEntry,
-    ConfigFlow,
-    ConfigFlowResult,
-    OptionsFlow,
-)
+from homeassistant.config_entries import ConfigFlow, ConfigFlowResult, OptionsFlow
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -41,9 +36,8 @@ from .const import (
     DEFAULT_SCAN_INTERVAL,
     DEFAULT_TELNET_PORT,
     DOMAIN,
-    ROUTER,
 )
-from .router import KeeneticRouter
+from .router import KeeneticConfigEntry
 
 
 class KeeneticFlowHandler(ConfigFlow, domain=DOMAIN):
@@ -56,7 +50,7 @@ class KeeneticFlowHandler(ConfigFlow, domain=DOMAIN):
     @staticmethod
     @callback
     def async_get_options_flow(
-        config_entry: ConfigEntry,
+        config_entry: KeeneticConfigEntry,
     ) -> KeeneticOptionsFlowHandler:
         """Get the options flow for this handler."""
         return KeeneticOptionsFlowHandler()
@@ -142,6 +136,8 @@ class KeeneticFlowHandler(ConfigFlow, domain=DOMAIN):
 class KeeneticOptionsFlowHandler(OptionsFlow):
     """Handle options."""
 
+    config_entry: KeeneticConfigEntry
+
     def __init__(self) -> None:
         """Initialize options flow."""
         self._interface_options: dict[str, str] = {}
@@ -150,9 +146,7 @@ class KeeneticOptionsFlowHandler(OptionsFlow):
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Manage the options."""
-        router: KeeneticRouter = self.hass.data[DOMAIN][self.config_entry.entry_id][
-            ROUTER
-        ]
+        router = self.config_entry.runtime_data
 
         interfaces: list[InterfaceInfo] = await self.hass.async_add_executor_job(
             router.client.get_interfaces

homeassistant/components/keenetic_ndms2/const.py

@@ -5,7 +5,6 @@ from homeassistant.components.device_tracker import (
 )
 
 DOMAIN = "keenetic_ndms2"
-ROUTER = "router"
 DEFAULT_TELNET_PORT = 23
 DEFAULT_SCAN_INTERVAL = 120
 DEFAULT_CONSIDER_HOME = _DEFAULT_CONSIDER_HOME.total_seconds()

homeassistant/components/keenetic_ndms2/device_tracker.py

@@ -10,26 +10,24 @@ from homeassistant.components.device_tracker import (
     DOMAIN as DEVICE_TRACKER_DOMAIN,
     ScannerEntity,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.util import dt as dt_util
 
-from .const import DOMAIN, ROUTER
-from .router import KeeneticRouter
+from .router import KeeneticConfigEntry, KeeneticRouter
 
 _LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: KeeneticConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up device tracker for Keenetic NDMS2 component."""
-    router: KeeneticRouter = hass.data[DOMAIN][config_entry.entry_id][ROUTER]
+    router = config_entry.runtime_data
 
     tracked: set[str] = set()
 

homeassistant/components/keenetic_ndms2/router.py

@@ -35,11 +35,13 @@ from .const import (
 
 _LOGGER = logging.getLogger(__name__)
 
+type KeeneticConfigEntry = ConfigEntry[KeeneticRouter]
+
 
 class KeeneticRouter:
     """Keenetic client Object."""
 
-    def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None:
+    def __init__(self, hass: HomeAssistant, config_entry: KeeneticConfigEntry) -> None:
         """Initialize the Client."""
         self.hass = hass
         self.config_entry = config_entry

tests/components/keenetic_ndms2/test_config_flow.py

@@ -87,19 +87,16 @@ async def test_options(hass: HomeAssistant) -> None:
     assert len(mock_setup_entry.mock_calls) == 1
 
     # fake router
-    hass.data.setdefault(keenetic.DOMAIN, {})
-    hass.data[keenetic.DOMAIN][entry.entry_id] = {
-        keenetic.ROUTER: Mock(
-            client=Mock(
-                get_interfaces=Mock(
-                    return_value=[
-                        InterfaceInfo.from_dict({"id": name, "type": "bridge"})
-                        for name in MOCK_OPTIONS[const.CONF_INTERFACES]
-                    ]
-                )
+    entry.runtime_data = Mock(
+        client=Mock(
+            get_interfaces=Mock(
+                return_value=[
+                    InterfaceInfo.from_dict({"id": name, "type": "bridge"})
+                    for name in MOCK_OPTIONS[const.CONF_INTERFACES]
+                ]
             )
         )
-    }
+    )
 
     result = await hass.config_entries.options.async_init(entry.entry_id)