Store runtime data inside the config entry in AVM Fritz!Smarthome (#116523)

homeassistant/components/fritzbox/__init__.py

@@ -4,52 +4,23 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 
-from pyfritzhome import Fritzhome, FritzhomeDevice, LoginError
+from pyfritzhome import FritzhomeDevice
 from pyfritzhome.devicetypes.fritzhomeentitybase import FritzhomeEntityBase
-from requests.exceptions import ConnectionError as RequestConnectionError
 
 from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR_DOMAIN
-from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import (
-    CONF_HOST,
-    CONF_PASSWORD,
-    CONF_USERNAME,
-    EVENT_HOMEASSISTANT_STOP,
-    UnitOfTemperature,
-)
+from homeassistant.const import EVENT_HOMEASSISTANT_STOP, UnitOfTemperature
 from homeassistant.core import Event, HomeAssistant
-from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers.device_registry import DeviceEntry, DeviceInfo
 from homeassistant.helpers.entity import EntityDescription
 from homeassistant.helpers.entity_registry import RegistryEntry, async_migrate_entries
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from .const import CONF_CONNECTIONS, CONF_COORDINATOR, DOMAIN, LOGGER, PLATFORMS
-from .coordinator import FritzboxDataUpdateCoordinator
+from .const import DOMAIN, LOGGER, PLATFORMS
+from .coordinator import FritzboxConfigEntry, FritzboxDataUpdateCoordinator
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: FritzboxConfigEntry) -> bool:
     """Set up the AVM FRITZ!SmartHome platforms."""
-    fritz = Fritzhome(
-        host=entry.data[CONF_HOST],
-        user=entry.data[CONF_USERNAME],
-        password=entry.data[CONF_PASSWORD],
-    )
-
-    try:
-        await hass.async_add_executor_job(fritz.login)
-    except RequestConnectionError as err:
-        raise ConfigEntryNotReady from err
-    except LoginError as err:
-        raise ConfigEntryAuthFailed from err
-
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = {
-        CONF_CONNECTIONS: fritz,
-    }
-
-    has_templates = await hass.async_add_executor_job(fritz.has_templates)
-    LOGGER.debug("enable smarthome templates: %s", has_templates)
 
     def _update_unique_id(entry: RegistryEntry) -> dict[str, str] | None:
         """Update unique ID of entity entry."""
@@ -73,15 +44,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     await async_migrate_entries(hass, entry.entry_id, _update_unique_id)
 
-    coordinator = FritzboxDataUpdateCoordinator(hass, entry.entry_id, has_templates)
+    coordinator = FritzboxDataUpdateCoordinator(hass, entry.entry_id)
     await coordinator.async_setup()
-    hass.data[DOMAIN][entry.entry_id][CONF_COORDINATOR] = coordinator
+
+    entry.runtime_data = coordinator
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     def logout_fritzbox(event: Event) -> None:
         """Close connections to this fritzbox."""
-        fritz.logout()
+        coordinator.fritz.logout()
 
     entry.async_on_unload(
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, logout_fritzbox)
@@ -90,25 +62,18 @@ 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: FritzboxConfigEntry) -> bool:
     """Unloading the AVM FRITZ!SmartHome platforms."""
-    fritz = hass.data[DOMAIN][entry.entry_id][CONF_CONNECTIONS]
-    await hass.async_add_executor_job(fritz.logout)
+    await hass.async_add_executor_job(entry.runtime_data.fritz.logout)
 
-    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)
 
 
 async def async_remove_config_entry_device(
-    hass: HomeAssistant, entry: ConfigEntry, device: DeviceEntry
+    hass: HomeAssistant, entry: FritzboxConfigEntry, device: DeviceEntry
 ) -> bool:
     """Remove Fritzbox config entry from a device."""
-    coordinator: FritzboxDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][
-        CONF_COORDINATOR
-    ]
+    coordinator = entry.runtime_data
 
     for identifier in device.identifiers:
         if identifier[0] == DOMAIN and (

homeassistant/components/fritzbox/binary_sensor.py

@@ -13,13 +13,12 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import FritzBoxDeviceEntity
-from .common import get_coordinator
+from .coordinator import FritzboxConfigEntry
 from .model import FritzEntityDescriptionMixinBase
 
 
@@ -65,10 +64,12 @@ BINARY_SENSOR_TYPES: Final[tuple[FritzBinarySensorEntityDescription, ...]] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: FritzboxConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the FRITZ!SmartHome binary sensor from ConfigEntry."""
-    coordinator = get_coordinator(hass, entry.entry_id)
+    coordinator = entry.runtime_data
 
     @callback
     def _add_entities(devices: set[str] | None = None) -> None:

homeassistant/components/fritzbox/button.py

@@ -3,21 +3,22 @@
 from pyfritzhome.devicetypes import FritzhomeTemplate
 
 from homeassistant.components.button import ButtonEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import FritzBoxEntity
-from .common import get_coordinator
 from .const import DOMAIN
+from .coordinator import FritzboxConfigEntry
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: FritzboxConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the FRITZ!SmartHome template from ConfigEntry."""
-    coordinator = get_coordinator(hass, entry.entry_id)
+    coordinator = entry.runtime_data
 
     @callback
     def _add_entities(templates: set[str] | None = None) -> None:

homeassistant/components/fritzbox/climate.py

@@ -12,7 +12,6 @@ from homeassistant.components.climate import (
     ClimateEntityFeature,
     HVACMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_BATTERY_LEVEL,
     ATTR_TEMPERATURE,
@@ -23,7 +22,6 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import FritzBoxDeviceEntity
-from .common import get_coordinator
 from .const import (
     ATTR_STATE_BATTERY_LOW,
     ATTR_STATE_HOLIDAY_MODE,
@@ -31,6 +29,7 @@ from .const import (
     ATTR_STATE_WINDOW_OPEN,
     LOGGER,
 )
+from .coordinator import FritzboxConfigEntry
 from .model import ClimateExtraAttributes
 
 OPERATION_LIST = [HVACMode.HEAT, HVACMode.OFF]
@@ -48,10 +47,12 @@ OFF_REPORT_SET_TEMPERATURE = 0.0
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: FritzboxConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the FRITZ!SmartHome thermostat from ConfigEntry."""
-    coordinator = get_coordinator(hass, entry.entry_id)
+    coordinator = entry.runtime_data
 
     @callback
     def _add_entities(devices: set[str] | None = None) -> None:

homeassistant/components/fritzbox/common.py

@@ -1,16 +0,0 @@
-"""Common functions for fritzbox integration."""
-
-from homeassistant.core import HomeAssistant
-
-from .const import CONF_COORDINATOR, DOMAIN
-from .coordinator import FritzboxDataUpdateCoordinator
-
-
-def get_coordinator(
-    hass: HomeAssistant, config_entry_id: str
-) -> FritzboxDataUpdateCoordinator:
-    """Get coordinator for given config entry id."""
-    coordinator: FritzboxDataUpdateCoordinator = hass.data[DOMAIN][config_entry_id][
-        CONF_COORDINATOR
-    ]
-    return coordinator

homeassistant/components/fritzbox/const.py

@@ -15,9 +15,6 @@ ATTR_STATE_WINDOW_OPEN: Final = "window_open"
 COLOR_MODE: Final = "1"
 COLOR_TEMP_MODE: Final = "4"
 
-CONF_CONNECTIONS: Final = "connections"
-CONF_COORDINATOR: Final = "coordinator"
-
 DEFAULT_HOST: Final = "fritz.box"
 DEFAULT_USERNAME: Final = "admin"
 

homeassistant/components/fritzbox/coordinator.py

@@ -10,12 +10,15 @@ from pyfritzhome.devicetypes import FritzhomeTemplate
 from requests.exceptions import ConnectionError as RequestConnectionError, HTTPError
 
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryAuthFailed
+from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import CONF_CONNECTIONS, DOMAIN, LOGGER
+from .const import DOMAIN, LOGGER
+
+FritzboxConfigEntry = ConfigEntry["FritzboxDataUpdateCoordinator"]
 
 
 @dataclass
@@ -29,10 +32,12 @@ class FritzboxCoordinatorData:
 class FritzboxDataUpdateCoordinator(DataUpdateCoordinator[FritzboxCoordinatorData]):
     """Fritzbox Smarthome device data update coordinator."""
 
-    config_entry: ConfigEntry
+    config_entry: FritzboxConfigEntry
     configuration_url: str
+    fritz: Fritzhome
+    has_templates: bool
 
-    def __init__(self, hass: HomeAssistant, name: str, has_templates: bool) -> None:
+    def __init__(self, hass: HomeAssistant, name: str) -> None:
         """Initialize the Fritzbox Smarthome device coordinator."""
         super().__init__(
             hass,
@@ -41,11 +46,6 @@ class FritzboxDataUpdateCoordinator(DataUpdateCoordinator[FritzboxCoordinatorDat
             update_interval=timedelta(seconds=30),
         )
 
-        self.fritz: Fritzhome = hass.data[DOMAIN][self.config_entry.entry_id][
-            CONF_CONNECTIONS
-        ]
-        self.configuration_url = self.fritz.get_prefixed_host()
-        self.has_templates = has_templates
         self.new_devices: set[str] = set()
         self.new_templates: set[str] = set()
 
@@ -53,6 +53,27 @@ class FritzboxDataUpdateCoordinator(DataUpdateCoordinator[FritzboxCoordinatorDat
 
     async def async_setup(self) -> None:
         """Set up the coordinator."""
+
+        self.fritz = Fritzhome(
+            host=self.config_entry.data[CONF_HOST],
+            user=self.config_entry.data[CONF_USERNAME],
+            password=self.config_entry.data[CONF_PASSWORD],
+        )
+
+        try:
+            await self.hass.async_add_executor_job(self.fritz.login)
+        except RequestConnectionError as err:
+            raise ConfigEntryNotReady from err
+        except LoginError as err:
+            raise ConfigEntryAuthFailed from err
+
+        self.has_templates = await self.hass.async_add_executor_job(
+            self.fritz.has_templates
+        )
+        LOGGER.debug("enable smarthome templates: %s", self.has_templates)
+
+        self.configuration_url = self.fritz.get_prefixed_host()
+
         await self.async_config_entry_first_refresh()
         self.cleanup_removed_devices(
             list(self.data.devices) + list(self.data.templates)

homeassistant/components/fritzbox/cover.py

@@ -10,19 +10,20 @@ from homeassistant.components.cover import (
     CoverEntity,
     CoverEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import FritzBoxDeviceEntity
-from .common import get_coordinator
+from .coordinator import FritzboxConfigEntry
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: FritzboxConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the FRITZ!SmartHome cover from ConfigEntry."""
-    coordinator = get_coordinator(hass, entry.entry_id)
+    coordinator = entry.runtime_data
 
     @callback
     def _add_entities(devices: set[str] | None = None) -> None:

homeassistant/components/fritzbox/diagnostics.py

@@ -5,22 +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 CONF_COORDINATOR, DOMAIN
-from .coordinator import FritzboxDataUpdateCoordinator
+from .coordinator import FritzboxConfigEntry
 
 TO_REDACT = {CONF_USERNAME, CONF_PASSWORD}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: FritzboxConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    data: dict = hass.data[DOMAIN][entry.entry_id]
-    coordinator: FritzboxDataUpdateCoordinator = data[CONF_COORDINATOR]
+    coordinator = entry.runtime_data
 
     diag_data = {
         "entry": async_redact_data(entry.as_dict(), TO_REDACT),

homeassistant/components/fritzbox/light.py

@@ -13,22 +13,23 @@ from homeassistant.components.light import (
     ColorMode,
     LightEntity,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import FritzboxDataUpdateCoordinator, FritzBoxDeviceEntity
-from .common import get_coordinator
 from .const import COLOR_MODE, COLOR_TEMP_MODE, LOGGER
+from .coordinator import FritzboxConfigEntry
 
 SUPPORTED_COLOR_MODES = {ColorMode.COLOR_TEMP, ColorMode.HS}
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: FritzboxConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the FRITZ!SmartHome light from ConfigEntry."""
-    coordinator = get_coordinator(hass, entry.entry_id)
+    coordinator = entry.runtime_data
 
     @callback
     def _add_entities(devices: set[str] | None = None) -> None:

homeassistant/components/fritzbox/sensor.py

@@ -16,7 +16,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     PERCENTAGE,
     EntityCategory,
@@ -32,7 +31,7 @@ from homeassistant.helpers.typing import StateType
 from homeassistant.util.dt import utc_from_timestamp
 
 from . import FritzBoxDeviceEntity
-from .common import get_coordinator
+from .coordinator import FritzboxConfigEntry
 from .model import FritzEntityDescriptionMixinBase
 
 
@@ -210,10 +209,12 @@ SENSOR_TYPES: Final[tuple[FritzSensorEntityDescription, ...]] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: FritzboxConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the FRITZ!SmartHome sensor from ConfigEntry."""
-    coordinator = get_coordinator(hass, entry.entry_id)
+    coordinator = entry.runtime_data
 
     @callback
     def _add_entities(devices: set[str] | None = None) -> None:

homeassistant/components/fritzbox/switch.py

@@ -5,19 +5,20 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.switch import SwitchEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import FritzBoxDeviceEntity
-from .common import get_coordinator
+from .coordinator import FritzboxConfigEntry
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: FritzboxConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the FRITZ!SmartHome switch from ConfigEntry."""
-    coordinator = get_coordinator(hass, entry.entry_id)
+    coordinator = entry.runtime_data
 
     @callback
     def _add_entities(devices: set[str] | None = None) -> None:

tests/components/fritzbox/conftest.py

@@ -9,7 +9,7 @@ import pytest
 def fritz_fixture() -> Mock:
     """Patch libraries."""
     with (
-        patch("homeassistant.components.fritzbox.Fritzhome") as fritz,
+        patch("homeassistant.components.fritzbox.coordinator.Fritzhome") as fritz,
         patch("homeassistant.components.fritzbox.config_flow.Fritzhome"),
     ):
         fritz.return_value.get_prefixed_host.return_value = "http://1.2.3.4"

tests/components/fritzbox/test_init.py

@@ -254,7 +254,7 @@ async def test_raise_config_entry_not_ready_when_offline(hass: HomeAssistant) ->
     )
     entry.add_to_hass(hass)
     with patch(
-        "homeassistant.components.fritzbox.Fritzhome.login",
+        "homeassistant.components.fritzbox.coordinator.Fritzhome.login",
         side_effect=RequestConnectionError(),
     ) as mock_login:
         await hass.config_entries.async_setup(entry.entry_id)
@@ -275,7 +275,7 @@ async def test_raise_config_entry_error_when_login_fail(hass: HomeAssistant) ->
     )
     entry.add_to_hass(hass)
     with patch(
-        "homeassistant.components.fritzbox.Fritzhome.login",
+        "homeassistant.components.fritzbox.coordinator.Fritzhome.login",
         side_effect=LoginError("user"),
     ) as mock_login:
         await hass.config_entries.async_setup(entry.entry_id)