Use ConfigEntry.runtime_data in AVM Fritz!Box tools (#136386)

* implement FritzConfigEntry with runtime_data

* use HassKey for platform global data

* update quality scale

* fix after rebase

* use FritzConfigEntry everywhere possible

* fix import of FritzConfigEntry in services.py

* pass the config_entry explicitly in coordinator init

* improve typing of FritzData

* use FritzConfigEntry in config_flow.py

homeassistant/components/fritz/__init__.py

@@ -2,7 +2,6 @@
 
 import logging
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -16,14 +15,13 @@ from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.typing import ConfigType
 
 from .const import (
-    DATA_FRITZ,
     DEFAULT_SSL,
     DOMAIN,
     FRITZ_AUTH_EXCEPTIONS,
     FRITZ_EXCEPTIONS,
     PLATFORMS,
 )
-from .coordinator import AvmWrapper, FritzData
+from .coordinator import FRITZ_DATA_KEY, AvmWrapper, FritzConfigEntry, FritzData
 from .services import async_setup_services
 
 _LOGGER = logging.getLogger(__name__)
@@ -37,11 +35,12 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: FritzConfigEntry) -> bool:
     """Set up fritzboxtools from config entry."""
     _LOGGER.debug("Setting up FRITZ!Box Tools component")
     avm_wrapper = AvmWrapper(
         hass=hass,
+        config_entry=entry,
         host=entry.data[CONF_HOST],
         port=entry.data[CONF_PORT],
         username=entry.data[CONF_USERNAME],
@@ -64,11 +63,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     await avm_wrapper.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})
+    entry.runtime_data = avm_wrapper
-    hass.data[DOMAIN][entry.entry_id] = avm_wrapper
 
-    if DATA_FRITZ not in hass.data:
+    if FRITZ_DATA_KEY not in hass.data:
-        hass.data[DATA_FRITZ] = FritzData()
+        hass.data[FRITZ_DATA_KEY] = FritzData()
 
     entry.async_on_unload(entry.add_update_listener(update_listener))
 
@@ -78,24 +76,20 @@ 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: FritzConfigEntry) -> bool:
     """Unload FRITZ!Box Tools config entry."""
-    avm_wrapper: AvmWrapper = hass.data[DOMAIN][entry.entry_id]
+    avm_wrapper = entry.runtime_data
 
-    fritz_data = hass.data[DATA_FRITZ]
+    fritz_data = hass.data[FRITZ_DATA_KEY]
     fritz_data.tracked.pop(avm_wrapper.unique_id)
 
     if not bool(fritz_data.tracked):
-        hass.data.pop(DATA_FRITZ)
+        hass.data.pop(FRITZ_DATA_KEY)
 
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        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: FritzConfigEntry) -> None:
     """Update when config_entry options update."""
     if entry.options:
         await hass.config_entries.async_reload(entry.entry_id)

homeassistant/components/fritz/binary_sensor.py

@@ -11,13 +11,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 AddEntitiesCallback
 
-from .const import DOMAIN
+from .coordinator import ConnectionInfo, FritzConfigEntry
-from .coordinator import AvmWrapper, ConnectionInfo
 from .entity import FritzBoxBaseCoordinatorEntity, FritzEntityDescription
 
 _LOGGER = logging.getLogger(__name__)
@@ -51,11 +49,13 @@ SENSOR_TYPES: tuple[FritzBinarySensorEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: FritzConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up entry."""
     _LOGGER.debug("Setting up FRITZ!Box binary sensors")
-    avm_wrapper: AvmWrapper = hass.data[DOMAIN][entry.entry_id]
+    avm_wrapper = entry.runtime_data
 
     connection_info = await avm_wrapper.async_get_connection_info()
 

homeassistant/components/fritz/button.py

@@ -12,15 +12,21 @@ from homeassistant.components.button import (
     ButtonEntity,
     ButtonEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import BUTTON_TYPE_WOL, CONNECTION_TYPE_LAN, DATA_FRITZ, DOMAIN, MeshRoles
+from .const import BUTTON_TYPE_WOL, CONNECTION_TYPE_LAN, DOMAIN, MeshRoles
-from .coordinator import AvmWrapper, FritzData, FritzDevice, _is_tracked
+from .coordinator import (
+    FRITZ_DATA_KEY,
+    AvmWrapper,
+    FritzConfigEntry,
+    FritzData,
+    FritzDevice,
+    _is_tracked,
+)
 from .entity import FritzDeviceBase
 
 _LOGGER = logging.getLogger(__name__)
@@ -65,12 +71,12 @@ BUTTONS: Final = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: FritzConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set buttons for device."""
     _LOGGER.debug("Setting up buttons")
-    avm_wrapper: AvmWrapper = hass.data[DOMAIN][entry.entry_id]
+    avm_wrapper = entry.runtime_data
 
     entities_list: list[ButtonEntity] = [
         FritzButton(avm_wrapper, entry.title, button) for button in BUTTONS
@@ -80,7 +86,7 @@ async def async_setup_entry(
         async_add_entities(entities_list)
         return
 
-    data_fritz: FritzData = hass.data[DATA_FRITZ]
+    data_fritz = hass.data[FRITZ_DATA_KEY]
     entities_list += _async_wol_buttons_list(avm_wrapper, data_fritz)
 
     async_add_entities(entities_list)

homeassistant/components/fritz/config_flow.py

@@ -17,12 +17,7 @@ from homeassistant.components.device_tracker import (
     CONF_CONSIDER_HOME,
     DEFAULT_CONSIDER_HOME,
 )
-from homeassistant.config_entries import (
+from homeassistant.config_entries import ConfigFlow, ConfigFlowResult, OptionsFlow
-    ConfigEntry,
-    ConfigFlow,
-    ConfigFlowResult,
-    OptionsFlow,
-)
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -53,6 +48,7 @@ from .const import (
     ERROR_UPNP_NOT_CONFIGURED,
     FRITZ_AUTH_EXCEPTIONS,
 )
+from .coordinator import FritzConfigEntry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -67,7 +63,7 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
     @staticmethod
     @callback
     def async_get_options_flow(
-        config_entry: ConfigEntry,
+        config_entry: FritzConfigEntry,
     ) -> FritzBoxToolsOptionsFlowHandler:
         """Get the options flow for this handler."""
         return FritzBoxToolsOptionsFlowHandler()
@@ -116,7 +112,7 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
 
         return None
 
-    async def async_check_configured_entry(self) -> ConfigEntry | None:
+    async def async_check_configured_entry(self) -> FritzConfigEntry | None:
         """Check if entry is configured."""
         current_host = await self.hass.async_add_executor_job(
             socket.gethostbyname, self._host

homeassistant/components/fritz/const.py

@@ -40,8 +40,6 @@ PLATFORMS = [
 CONF_OLD_DISCOVERY = "old_discovery"
 DEFAULT_CONF_OLD_DISCOVERY = False
 
-DATA_FRITZ = "fritz_data"
-
 DSL_CONNECTION: Literal["dsl"] = "dsl"
 
 DEFAULT_DEVICE_NAME = "Unknown device"

homeassistant/components/fritz/coordinator.py

@@ -36,6 +36,7 @@ from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 from homeassistant.util import dt as dt_util
+from homeassistant.util.hass_dict import HassKey
 
 from .const import (
     CONF_OLD_DISCOVERY,
@@ -50,8 +51,12 @@ from .const import (
 
 _LOGGER = logging.getLogger(__name__)
 
+FRITZ_DATA_KEY: HassKey[FritzData] = HassKey(DOMAIN)
 
-def _is_tracked(mac: str, current_devices: ValuesView) -> bool:
+type FritzConfigEntry = ConfigEntry[AvmWrapper]
+
+
+def _is_tracked(mac: str, current_devices: ValuesView[set[str]]) -> bool:
     """Check if device is already tracked."""
     return any(mac in tracked for tracked in current_devices)
 
@@ -59,7 +64,7 @@ def _is_tracked(mac: str, current_devices: ValuesView) -> bool:
 def device_filter_out_from_trackers(
     mac: str,
     device: FritzDevice,
-    current_devices: ValuesView,
+    current_devices: ValuesView[set[str]],
 ) -> bool:
     """Check if device should be filtered out from trackers."""
     reason: str | None = None
@@ -160,11 +165,12 @@ class UpdateCoordinatorDataType(TypedDict):
 class FritzBoxTools(DataUpdateCoordinator[UpdateCoordinatorDataType]):
     """FritzBoxTools class."""
 
-    config_entry: ConfigEntry
+    config_entry: FritzConfigEntry
 
     def __init__(
         self,
         hass: HomeAssistant,
+        config_entry: FritzConfigEntry,
         password: str,
         port: int,
         username: str = DEFAULT_USERNAME,
@@ -174,6 +180,7 @@ class FritzBoxTools(DataUpdateCoordinator[UpdateCoordinatorDataType]):
         """Initialize FritzboxTools class."""
         super().__init__(
             hass=hass,
+            config_entry=config_entry,
             logger=_LOGGER,
             name=f"{DOMAIN}-{host}-coordinator",
             update_interval=timedelta(seconds=30),
@@ -869,9 +876,9 @@ class AvmWrapper(FritzBoxTools):
 class FritzData:
     """Storage class for platform global data."""
 
-    tracked: dict = field(default_factory=dict)
+    tracked: dict[str, set[str]] = field(default_factory=dict)
-    profile_switches: dict = field(default_factory=dict)
+    profile_switches: dict[str, set[str]] = field(default_factory=dict)
-    wol_buttons: dict = field(default_factory=dict)
+    wol_buttons: dict[str, set[str]] = field(default_factory=dict)
 
 
 class FritzDevice:

homeassistant/components/fritz/device_tracker.py

@@ -6,14 +6,14 @@ import datetime
 import logging
 
 from homeassistant.components.device_tracker import ScannerEntity
-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_FRITZ, DOMAIN
 from .coordinator import (
+    FRITZ_DATA_KEY,
     AvmWrapper,
+    FritzConfigEntry,
     FritzData,
     FritzDevice,
     device_filter_out_from_trackers,
@@ -24,12 +24,14 @@ _LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: FritzConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up device tracker for FRITZ!Box component."""
     _LOGGER.debug("Starting FRITZ!Box device tracker")
-    avm_wrapper: AvmWrapper = hass.data[DOMAIN][entry.entry_id]
+    avm_wrapper = entry.runtime_data
-    data_fritz: FritzData = hass.data[DATA_FRITZ]
+    data_fritz = hass.data[FRITZ_DATA_KEY]
 
     @callback
     def update_avm_device() -> None:

homeassistant/components/fritz/diagnostics.py

@@ -5,21 +5,19 @@ 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_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
+from .coordinator import FritzConfigEntry
-from .coordinator import AvmWrapper
 
 TO_REDACT = {CONF_USERNAME, CONF_PASSWORD}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: FritzConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    avm_wrapper: AvmWrapper = hass.data[DOMAIN][entry.entry_id]
+    avm_wrapper = entry.runtime_data
 
     return {
         "entry": async_redact_data(entry.as_dict(), TO_REDACT),

homeassistant/components/fritz/image.py

@@ -8,14 +8,12 @@ import logging
 from requests.exceptions import RequestException
 
 from homeassistant.components.image import ImageEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import dt as dt_util, slugify
 
-from .const import DOMAIN
+from .coordinator import AvmWrapper, FritzConfigEntry
-from .coordinator import AvmWrapper
 from .entity import FritzBoxBaseEntity
 
 _LOGGER = logging.getLogger(__name__)
@@ -23,11 +21,11 @@ _LOGGER = logging.getLogger(__name__)
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: FritzConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up guest WiFi QR code for device."""
-    avm_wrapper: AvmWrapper = hass.data[DOMAIN][entry.entry_id]
+    avm_wrapper = entry.runtime_data
 
     guest_wifi_info = await hass.async_add_executor_job(
         avm_wrapper.fritz_guest_wifi.get_info

homeassistant/components/fritz/quality_scale.yaml

@@ -22,9 +22,7 @@ rules:
   has-entity-name:
     status: todo
     comment: partially done
-  runtime-data:
+  runtime-data: done
-    status: todo
-    comment: still uses hass.data
   test-before-configure: done
   test-before-setup: done
   unique-config-entry: done

homeassistant/components/fritz/sensor.py

@@ -15,7 +15,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     SIGNAL_STRENGTH_DECIBELS,
     EntityCategory,
@@ -27,8 +26,8 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.util.dt import utcnow
 
-from .const import DOMAIN, DSL_CONNECTION, UPTIME_DEVIATION
+from .const import DSL_CONNECTION, UPTIME_DEVIATION
-from .coordinator import AvmWrapper, ConnectionInfo
+from .coordinator import ConnectionInfo, FritzConfigEntry
 from .entity import FritzBoxBaseCoordinatorEntity, FritzEntityDescription
 
 _LOGGER = logging.getLogger(__name__)
@@ -267,11 +266,13 @@ SENSOR_TYPES: tuple[FritzSensorEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: FritzConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up entry."""
     _LOGGER.debug("Setting up FRITZ!Box sensors")
-    avm_wrapper: AvmWrapper = hass.data[DOMAIN][entry.entry_id]
+    avm_wrapper = entry.runtime_data
 
     connection_info = await avm_wrapper.async_get_connection_info()
 

homeassistant/components/fritz/services.py

@@ -15,7 +15,7 @@ from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
 from homeassistant.helpers.service import async_extract_config_entry_ids
 
 from .const import DOMAIN
-from .coordinator import AvmWrapper
+from .coordinator import FritzConfigEntry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -33,7 +33,7 @@ async def _async_set_guest_wifi_password(service_call: ServiceCall) -> None:
     """Call Fritz set guest wifi password service."""
     hass = service_call.hass
     target_entry_ids = await async_extract_config_entry_ids(hass, service_call)
-    target_entries = [
+    target_entries: list[FritzConfigEntry] = [
         loaded_entry
         for loaded_entry in hass.config_entries.async_loaded_entries(DOMAIN)
         if loaded_entry.entry_id in target_entry_ids
@@ -48,7 +48,7 @@ async def _async_set_guest_wifi_password(service_call: ServiceCall) -> None:
 
     for target_entry in target_entries:
         _LOGGER.debug("Executing service %s", service_call.service)
-        avm_wrapper: AvmWrapper = hass.data[DOMAIN][target_entry.entry_id]
+        avm_wrapper = target_entry.runtime_data
         try:
             await avm_wrapper.async_trigger_set_guest_password(
                 service_call.data.get("password"),

homeassistant/components/fritz/switch.py

@@ -7,7 +7,6 @@ from typing import Any
 
 from homeassistant.components.network import async_get_source_ip
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo
@@ -18,7 +17,6 @@ from homeassistant.helpers.update_coordinator import CoordinatorEntity
 from homeassistant.util import slugify
 
 from .const import (
-    DATA_FRITZ,
     DOMAIN,
     SWITCH_TYPE_DEFLECTION,
     SWITCH_TYPE_PORTFORWARD,
@@ -28,7 +26,9 @@ from .const import (
     MeshRoles,
 )
 from .coordinator import (
+    FRITZ_DATA_KEY,
     AvmWrapper,
+    FritzConfigEntry,
     FritzData,
     FritzDevice,
     SwitchInfo,
@@ -220,12 +220,14 @@ async def async_all_entities_list(
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: FritzConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up entry."""
     _LOGGER.debug("Setting up switches")
-    avm_wrapper: AvmWrapper = hass.data[DOMAIN][entry.entry_id]
+    avm_wrapper = entry.runtime_data
-    data_fritz: FritzData = hass.data[DATA_FRITZ]
+    data_fritz = hass.data[FRITZ_DATA_KEY]
 
     _LOGGER.debug("Fritzbox services: %s", avm_wrapper.connection.services)
 

homeassistant/components/fritz/update.py

@@ -11,13 +11,11 @@ from homeassistant.components.update import (
     UpdateEntityDescription,
     UpdateEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from .coordinator import AvmWrapper, FritzConfigEntry
-from .coordinator import AvmWrapper
 from .entity import FritzBoxBaseCoordinatorEntity, FritzEntityDescription
 
 _LOGGER = logging.getLogger(__name__)
@@ -29,11 +27,13 @@ class FritzUpdateEntityDescription(UpdateEntityDescription, FritzEntityDescripti
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: FritzConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up AVM FRITZ!Box update entities."""
     _LOGGER.debug("Setting up AVM FRITZ!Box update entities")
-    avm_wrapper: AvmWrapper = hass.data[DOMAIN][entry.entry_id]
+    avm_wrapper = entry.runtime_data
 
     entities = [FritzBoxUpdateEntity(avm_wrapper, entry.title)]