Use ConfigEntry.runtime_data to store runtime data at Home Connect (#131014)

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

homeassistant/components/home_connect/__init__.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from datetime import timedelta
 import logging
-from typing import Any
+from typing import Any, cast
 
 from requests import HTTPError
 import voluptuous as vol
@@ -40,6 +40,8 @@ from .const import (
     SERVICE_START_PROGRAM,
 )
 
+type HomeConnectConfigEntry = ConfigEntry[api.ConfigEntryAuth]
+
 _LOGGER = logging.getLogger(__name__)
 
 SCAN_INTERVAL = timedelta(minutes=1)
@@ -89,13 +91,17 @@ PLATFORMS = [
 ]
 
 
-def _get_appliance_by_device_id(
+def _get_appliance(
-    hass: HomeAssistant, device_id: str
+    hass: HomeAssistant,
+    device_id: str | None = None,
+    device_entry: dr.DeviceEntry | None = None,
+    entry: HomeConnectConfigEntry | None = None,
 ) -> api.HomeConnectAppliance:
-    """Return a Home Connect appliance instance given an device_id."""
+    """Return a Home Connect appliance instance given a device id or a device entry."""
+    if device_id is not None and device_entry is None:
         device_registry = dr.async_get(hass)
         device_entry = device_registry.async_get(device_id)
-    assert device_entry
+    assert device_entry, "Either a device id or a device entry must be provided"
 
     ha_id = next(
         (
@@ -107,17 +113,30 @@ def _get_appliance_by_device_id(
     )
     assert ha_id
 
-    for hc_api in hass.data[DOMAIN].values():
+    def find_appliance(
-        for device in hc_api.devices:
+        entry: HomeConnectConfigEntry,
+    ) -> api.HomeConnectAppliance | None:
+        for device in entry.runtime_data.devices:
             appliance = device.appliance
             if appliance.haId == ha_id:
                 return appliance
-    raise ValueError(f"Appliance for device id {device_id} not found")
+        return None
+
+    if entry is None:
+        for entry_id in device_entry.config_entries:
+            entry = hass.config_entries.async_get_entry(entry_id)
+            assert entry
+            if entry.domain == DOMAIN:
+                entry = cast(HomeConnectConfigEntry, entry)
+                if (appliance := find_appliance(entry)) is not None:
+                    return appliance
+    elif (appliance := find_appliance(entry)) is not None:
+        return appliance
+    raise ValueError(f"Appliance for device id {device_entry.id} not found")
 
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up Home Connect component."""
-    hass.data[DOMAIN] = {}
 
     async def _async_service_program(call, method):
         """Execute calls to services taking a program."""
@@ -136,14 +155,14 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
 
             options.append(option)
 
-        appliance = _get_appliance_by_device_id(hass, device_id)
+        appliance = _get_appliance(hass, device_id)
         await hass.async_add_executor_job(getattr(appliance, method), program, options)
 
     async def _async_service_command(call, command):
         """Execute calls to services executing a command."""
         device_id = call.data[ATTR_DEVICE_ID]
 
-        appliance = _get_appliance_by_device_id(hass, device_id)
+        appliance = _get_appliance(hass, device_id)
         await hass.async_add_executor_job(appliance.execute_command, command)
 
     async def _async_service_key_value(call, method):
@@ -153,7 +172,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
         unit = call.data.get(ATTR_UNIT)
         device_id = call.data[ATTR_DEVICE_ID]
 
-        appliance = _get_appliance_by_device_id(hass, device_id)
+        appliance = _get_appliance(hass, device_id)
         if unit is not None:
             await hass.async_add_executor_job(
                 getattr(appliance, method),
@@ -239,7 +258,7 @@ 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: HomeConnectConfigEntry) -> bool:
     """Set up Home Connect from a config entry."""
     implementation = (
         await config_entry_oauth2_flow.async_get_config_entry_implementation(
@@ -247,9 +266,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         )
     )
 
-    hc_api = api.ConfigEntryAuth(hass, entry, implementation)
+    entry.runtime_data = api.ConfigEntryAuth(hass, entry, implementation)
-
-    hass.data[DOMAIN][entry.entry_id] = hc_api
 
     await update_all_devices(hass, entry)
 
@@ -258,20 +275,19 @@ 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: HomeConnectConfigEntry
+) -> bool:
     """Unload a config entry."""
-    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
 
 
 @Throttle(SCAN_INTERVAL)
-async def update_all_devices(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def update_all_devices(
+    hass: HomeAssistant, entry: HomeConnectConfigEntry
+) -> None:
     """Update all the devices."""
-    data = hass.data[DOMAIN]
+    hc_api = entry.runtime_data
-    hc_api = data[entry.entry_id]
 
     try:
         await hass.async_add_executor_job(hc_api.get_devices)
@@ -281,11 +297,13 @@ async def update_all_devices(hass: HomeAssistant, entry: ConfigEntry) -> None:
         _LOGGER.warning("Cannot update devices: %s", err.response.status_code)
 
 
-async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_migrate_entry(
+    hass: HomeAssistant, entry: HomeConnectConfigEntry
+) -> bool:
     """Migrate old entry."""
-    _LOGGER.debug("Migrating from version %s", config_entry.version)
+    _LOGGER.debug("Migrating from version %s", entry.version)
 
-    if config_entry.version == 1 and config_entry.minor_version == 1:
+    if entry.version == 1 and entry.minor_version == 1:
 
         @callback
         def update_unique_id(
@@ -301,11 +319,11 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
                     }
             return None
 
-        await async_migrate_entries(hass, config_entry.entry_id, update_unique_id)
+        await async_migrate_entries(hass, entry.entry_id, update_unique_id)
 
-        hass.config_entries.async_update_entry(config_entry, minor_version=2)
+        hass.config_entries.async_update_entry(entry, minor_version=2)
 
-    _LOGGER.debug("Migration to version %s successful", config_entry.version)
+    _LOGGER.debug("Migration to version %s successful", entry.version)
     return True
 
 

homeassistant/components/home_connect/binary_sensor.py

@@ -10,7 +10,6 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntityDescription,
 )
 from homeassistant.components.script import scripts_with_entity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -20,6 +19,7 @@ from homeassistant.helpers.issue_registry import (
     async_delete_issue,
 )
 
+from . import HomeConnectConfigEntry
 from .api import HomeConnectDevice
 from .const import (
     ATTR_VALUE,
@@ -118,15 +118,14 @@ BINARY_SENSORS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: HomeConnectConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Home Connect binary sensor."""
 
     def get_entities() -> list[BinarySensorEntity]:
         entities: list[BinarySensorEntity] = []
-        hc_api = hass.data[DOMAIN][config_entry.entry_id]
+        for device in entry.runtime_data.devices:
-        for device in hc_api.devices:
             entities.extend(
                 HomeConnectBinarySensor(device, description)
                 for description in BINARY_SENSORS

homeassistant/components/home_connect/diagnostics.py

@@ -6,13 +6,11 @@ from typing import Any
 
 from homeconnect.api import HomeConnectAppliance
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceEntry
 
-from . import _get_appliance_by_device_id
+from . import HomeConnectConfigEntry, _get_appliance
 from .api import HomeConnectDevice
-from .const import DOMAIN
 
 
 def _generate_appliance_diagnostics(appliance: HomeConnectAppliance) -> dict[str, Any]:
@@ -32,17 +30,17 @@ def _generate_entry_diagnostics(
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, entry: HomeConnectConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
     return await hass.async_add_executor_job(
-        _generate_entry_diagnostics, hass.data[DOMAIN][config_entry.entry_id].devices
+        _generate_entry_diagnostics, entry.runtime_data.devices
     )
 
 
 async def async_get_device_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry, device: DeviceEntry
+    hass: HomeAssistant, entry: HomeConnectConfigEntry, device: DeviceEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a device."""
-    appliance = _get_appliance_by_device_id(hass, device.id)
+    appliance = _get_appliance(hass, device_entry=device, entry=entry)
     return await hass.async_add_executor_job(_generate_appliance_diagnostics, appliance)

homeassistant/components/home_connect/light.py

@@ -15,14 +15,13 @@ from homeassistant.components.light import (
     LightEntity,
     LightEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ServiceValidationError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.util.color as color_util
 
-from . import get_dict_from_home_connect_error
+from . import HomeConnectConfigEntry, get_dict_from_home_connect_error
-from .api import ConfigEntryAuth, HomeConnectDevice
+from .api import HomeConnectDevice
 from .const import (
     ATTR_VALUE,
     BSH_AMBIENT_LIGHT_BRIGHTNESS,
@@ -88,18 +87,17 @@ LIGHTS: tuple[HomeConnectLightEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: HomeConnectConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Home Connect light."""
 
     def get_entities() -> list[LightEntity]:
         """Get a list of entities."""
-        hc_api: ConfigEntryAuth = hass.data[DOMAIN][config_entry.entry_id]
         return [
             HomeConnectLight(device, description)
             for description in LIGHTS
-            for device in hc_api.devices
+            for device in entry.runtime_data.devices
             if description.key in device.appliance.status
         ]
 

homeassistant/components/home_connect/number.py

@@ -11,13 +11,11 @@ from homeassistant.components.number import (
     NumberEntity,
     NumberEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ServiceValidationError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import get_dict_from_home_connect_error
+from . import HomeConnectConfigEntry, get_dict_from_home_connect_error
-from .api import ConfigEntryAuth
 from .const import (
     ATTR_CONSTRAINTS,
     ATTR_STEPSIZE,
@@ -84,18 +82,17 @@ NUMBERS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: HomeConnectConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Home Connect number."""
 
     def get_entities() -> list[HomeConnectNumberEntity]:
         """Get a list of entities."""
-        hc_api: ConfigEntryAuth = hass.data[DOMAIN][config_entry.entry_id]
         return [
             HomeConnectNumberEntity(device, description)
             for description in NUMBERS
-            for device in hc_api.devices
+            for device in entry.runtime_data.devices
             if description.key in device.appliance.status
         ]
 

homeassistant/components/home_connect/sensor.py

@@ -14,14 +14,13 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE, UnitOfTime, UnitOfVolume
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 import homeassistant.util.dt as dt_util
 
-from .api import ConfigEntryAuth
+from . import HomeConnectConfigEntry
 from .const import (
     ATTR_VALUE,
     BSH_DOOR_STATE,
@@ -34,7 +33,6 @@ from .const import (
     COFFEE_EVENT_WATER_TANK_EMPTY,
     DISHWASHER_EVENT_RINSE_AID_NEARLY_EMPTY,
     DISHWASHER_EVENT_SALT_NEARLY_EMPTY,
-    DOMAIN,
     REFRIGERATION_EVENT_DOOR_ALARM_FREEZER,
     REFRIGERATION_EVENT_DOOR_ALARM_REFRIGERATOR,
     REFRIGERATION_EVENT_TEMP_ALARM_FREEZER,
@@ -253,7 +251,7 @@ EVENT_SENSORS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: HomeConnectConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Home Connect sensor."""
@@ -261,8 +259,7 @@ async def async_setup_entry(
     def get_entities() -> list[SensorEntity]:
         """Get a list of entities."""
         entities: list[SensorEntity] = []
-        hc_api: ConfigEntryAuth = hass.data[DOMAIN][config_entry.entry_id]
+        for device in entry.runtime_data.devices:
-        for device in hc_api.devices:
             entities.extend(
                 HomeConnectSensor(
                     device,

homeassistant/components/home_connect/switch.py

@@ -7,13 +7,11 @@ from typing import Any
 from homeconnect.api import HomeConnectError
 
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ServiceValidationError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import get_dict_from_home_connect_error
+from . import HomeConnectConfigEntry, get_dict_from_home_connect_error
-from .api import ConfigEntryAuth
 from .const import (
     ATTR_ALLOWED_VALUES,
     ATTR_CONSTRAINTS,
@@ -105,7 +103,7 @@ SWITCHES = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: HomeConnectConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Home Connect switch."""
@@ -113,8 +111,7 @@ async def async_setup_entry(
     def get_entities() -> list[SwitchEntity]:
         """Get a list of entities."""
         entities: list[SwitchEntity] = []
-        hc_api: ConfigEntryAuth = hass.data[DOMAIN][config_entry.entry_id]
+        for device in entry.runtime_data.devices:
-        for device in hc_api.devices:
             if device.appliance.type in APPLIANCES_WITH_PROGRAMS:
                 with contextlib.suppress(HomeConnectError):
                     programs = device.appliance.get_programs_available()

homeassistant/components/home_connect/time.py

@@ -6,13 +6,11 @@ import logging
 from homeconnect.api import HomeConnectError
 
 from homeassistant.components.time import TimeEntity, TimeEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ServiceValidationError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import get_dict_from_home_connect_error
+from . import HomeConnectConfigEntry, get_dict_from_home_connect_error
-from .api import ConfigEntryAuth
 from .const import (
     ATTR_VALUE,
     DOMAIN,
@@ -35,18 +33,17 @@ TIME_ENTITIES = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: HomeConnectConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Home Connect switch."""
 
     def get_entities() -> list[HomeConnectTimeEntity]:
         """Get a list of entities."""
-        hc_api: ConfigEntryAuth = hass.data[DOMAIN][config_entry.entry_id]
         return [
             HomeConnectTimeEntity(device, description)
             for description in TIME_ENTITIES
-            for device in hc_api.devices
+            for device in entry.runtime_data.devices
             if description.key in device.appliance.status
         ]
 

tests/components/home_connect/test_diagnostics.py

@@ -60,3 +60,28 @@ async def test_async_get_device_diagnostics(
     )
 
     assert await async_get_device_diagnostics(hass, config_entry, device) == snapshot
+
+
+@pytest.mark.usefixtures("bypass_throttle")
+async def test_async_device_diagnostics_exceptions(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    integration_setup: Callable[[], Awaitable[bool]],
+    setup_credentials: None,
+    get_appliances: MagicMock,
+    device_registry: dr.DeviceRegistry,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test device config entry diagnostics."""
+    get_appliances.side_effect = get_all_appliances
+    assert config_entry.state == ConfigEntryState.NOT_LOADED
+    assert await integration_setup()
+    assert config_entry.state == ConfigEntryState.LOADED
+
+    device = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={(DOMAIN, "Random-Device-ID")},
+    )
+
+    with pytest.raises(ValueError):
+        await async_get_device_diagnostics(hass, config_entry, device)