Use runtime_data in hue (#144946)

* Use runtime_data in hue

* More

* Tests

homeassistant/components/hue/__init__.py

@@ -3,17 +3,17 @@
 from aiohue.util import normalize_bridge_id
 
 from homeassistant.components import persistent_notification
-from homeassistant.config_entries import SOURCE_IGNORE, ConfigEntry
+from homeassistant.config_entries import SOURCE_IGNORE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 
-from .bridge import HueBridge
+from .bridge import HueBridge, HueConfigEntry
 from .const import DOMAIN, SERVICE_HUE_ACTIVATE_SCENE
 from .migration import check_migration
 from .services import async_register_services
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: HueConfigEntry) -> bool:
     """Set up a bridge from a config entry."""
     # check (and run) migrations if needed
     await check_migration(hass, entry)
@@ -104,10 +104,9 @@ 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: HueConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_success = await hass.data[DOMAIN][entry.entry_id].async_reset()
-    if len(hass.data[DOMAIN]) == 0:
-        hass.data.pop(DOMAIN)
+    unload_success = await entry.runtime_data.async_reset()
+    if not hass.config_entries.async_loaded_entries(DOMAIN):
         hass.services.async_remove(DOMAIN, SERVICE_HUE_ACTIVATE_SCENE)
     return unload_success

homeassistant/components/hue/binary_sensor.py

@@ -2,23 +2,21 @@
 
 from __future__ import annotations
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .bridge import HueBridge
-from .const import DOMAIN
+from .bridge import HueConfigEntry
 from .v1.binary_sensor import async_setup_entry as setup_entry_v1
 from .v2.binary_sensor import async_setup_entry as setup_entry_v2
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HueConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up binary sensor entities."""
-    bridge: HueBridge = hass.data[DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
     if bridge.api_version == 1:
         await setup_entry_v1(hass, config_entry, async_add_entities)
     else:

homeassistant/components/hue/bridge.py

@@ -36,11 +36,13 @@ PLATFORMS_v2 = [
     Platform.SWITCH,
 ]
 
+type HueConfigEntry = ConfigEntry[HueBridge]
+
 
 class HueBridge:
     """Manages a single Hue bridge."""
 
-    def __init__(self, hass: core.HomeAssistant, config_entry: ConfigEntry) -> None:
+    def __init__(self, hass: core.HomeAssistant, config_entry: HueConfigEntry) -> None:
         """Initialize the system."""
         self.config_entry = config_entry
         self.hass = hass
@@ -58,7 +60,7 @@ class HueBridge:
         else:
             self.api = HueBridgeV2(self.host, app_key)
         # store (this) bridge object in hass data
-        hass.data.setdefault(DOMAIN, {})[self.config_entry.entry_id] = self
+        self.config_entry.runtime_data = self
 
     @property
     def host(self) -> str:
@@ -163,7 +165,7 @@ class HueBridge:
         )
 
         if unload_success:
-            self.hass.data[DOMAIN].pop(self.config_entry.entry_id)
+            delattr(self.config_entry, "runtime_data")
 
         return unload_success
 
@@ -179,7 +181,7 @@ class HueBridge:
         create_config_flow(self.hass, self.host)
 
 
-async def _update_listener(hass: core.HomeAssistant, entry: ConfigEntry) -> None:
+async def _update_listener(hass: core.HomeAssistant, entry: HueConfigEntry) -> None:
     """Handle ConfigEntry options update."""
     await hass.config_entries.async_reload(entry.entry_id)
 

homeassistant/components/hue/config_flow.py

@@ -13,12 +13,7 @@ from aiohue.util import normalize_bridge_id
 import slugify as unicode_slug
 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_API_KEY, CONF_API_VERSION, CONF_HOST
 from homeassistant.core import callback
 from homeassistant.helpers import (
@@ -28,6 +23,7 @@ from homeassistant.helpers import (
 )
 from homeassistant.helpers.service_info.zeroconf import ZeroconfServiceInfo
 
+from .bridge import HueConfigEntry
 from .const import (
     CONF_ALLOW_HUE_GROUPS,
     CONF_ALLOW_UNREACHABLE,
@@ -53,7 +49,7 @@ class HueFlowHandler(ConfigFlow, domain=DOMAIN):
     @staticmethod
     @callback
     def async_get_options_flow(
-        config_entry: ConfigEntry,
+        config_entry: HueConfigEntry,
     ) -> HueV1OptionsFlowHandler | HueV2OptionsFlowHandler:
         """Get the options flow for this handler."""
         if config_entry.data.get(CONF_API_VERSION, 1) == 1:

homeassistant/components/hue/device_trigger.py

@@ -26,14 +26,15 @@ if TYPE_CHECKING:
     from homeassistant.core import HomeAssistant
     from homeassistant.helpers.trigger import TriggerActionType, TriggerInfo
 
-    from .bridge import HueBridge
+    from .bridge import HueConfigEntry
 
 
 async def async_validate_trigger_config(
     hass: HomeAssistant, config: ConfigType
 ) -> ConfigType:
     """Validate config."""
-    if DOMAIN not in hass.data:
+    entries: list[HueConfigEntry] = hass.config_entries.async_loaded_entries(DOMAIN)
+    if not entries:
         # happens at startup
         return config
     device_id = config[CONF_DEVICE_ID]
@@ -42,10 +43,10 @@ async def async_validate_trigger_config(
     if (device_entry := dev_reg.async_get(device_id)) is None:
         raise InvalidDeviceAutomationConfig(f"Device ID {device_id} is not valid")
 
-    for conf_entry_id in device_entry.config_entries:
-        if conf_entry_id not in hass.data[DOMAIN]:
+    for entry in entries:
+        if entry.entry_id not in device_entry.config_entries:
             continue
-        bridge: HueBridge = hass.data[DOMAIN][conf_entry_id]
+        bridge = entry.runtime_data
         if bridge.api_version == 1:
             return await async_validate_trigger_config_v1(bridge, device_entry, config)
         return await async_validate_trigger_config_v2(bridge, device_entry, config)
@@ -65,10 +66,11 @@ async def async_attach_trigger(
     if (device_entry := dev_reg.async_get(device_id)) is None:
         raise InvalidDeviceAutomationConfig(f"Device ID {device_id} is not valid")
 
-    for conf_entry_id in device_entry.config_entries:
-        if conf_entry_id not in hass.data[DOMAIN]:
+    entry: HueConfigEntry
+    for entry in hass.config_entries.async_loaded_entries(DOMAIN):
+        if entry.entry_id not in device_entry.config_entries:
             continue
-        bridge: HueBridge = hass.data[DOMAIN][conf_entry_id]
+        bridge = entry.runtime_data
         if bridge.api_version == 1:
             return await async_attach_trigger_v1(
                 bridge, device_entry, config, action, trigger_info
@@ -85,7 +87,8 @@ async def async_get_triggers(
     hass: HomeAssistant, device_id: str
 ) -> list[dict[str, Any]]:
     """Get device triggers for given (hass) device id."""
-    if DOMAIN not in hass.data:
+    entries: list[HueConfigEntry] = hass.config_entries.async_loaded_entries(DOMAIN)
+    if not entries:
         return []
     # lookup device in HASS DeviceRegistry
     dev_reg: dr.DeviceRegistry = dr.async_get(hass)
@@ -94,10 +97,10 @@ async def async_get_triggers(
 
     # Iterate all config entries for this device
     # and work out the bridge version
-    for conf_entry_id in device_entry.config_entries:
-        if conf_entry_id not in hass.data[DOMAIN]:
+    for entry in entries:
+        if entry.entry_id not in device_entry.config_entries:
             continue
-        bridge: HueBridge = hass.data[DOMAIN][conf_entry_id]
+        bridge = entry.runtime_data
 
         if bridge.api_version == 1:
             return async_get_triggers_v1(bridge, device_entry)

homeassistant/components/hue/diagnostics.py

@@ -4,18 +4,16 @@ from __future__ import annotations
 
 from typing import Any
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from .bridge import HueBridge
-from .const import DOMAIN
+from .bridge import HueConfigEntry
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: HueConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    bridge: HueBridge = hass.data[DOMAIN][entry.entry_id]
+    bridge = entry.runtime_data
     if bridge.api_version == 1:
         # diagnostics is only implemented for V2 bridges.
         return {}

homeassistant/components/hue/event.py

@@ -14,22 +14,21 @@ from homeassistant.components.event import (
     EventEntity,
     EventEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .bridge import HueBridge
-from .const import DEFAULT_BUTTON_EVENT_TYPES, DEVICE_SPECIFIC_EVENT_TYPES, DOMAIN
+from .bridge import HueConfigEntry
+from .const import DEFAULT_BUTTON_EVENT_TYPES, DEVICE_SPECIFIC_EVENT_TYPES
 from .v2.entity import HueBaseEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HueConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up event platform from Hue button resources."""
-    bridge: HueBridge = hass.data[DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
     api: HueBridgeV2 = bridge.api
 
     if bridge.api_version == 1:

homeassistant/components/hue/light.py

@@ -2,12 +2,10 @@
 
 from __future__ import annotations
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .bridge import HueBridge
-from .const import DOMAIN
+from .bridge import HueConfigEntry
 from .v1.light import async_setup_entry as setup_entry_v1
 from .v2.group import async_setup_entry as setup_groups_entry_v2
 from .v2.light import async_setup_entry as setup_entry_v2
@@ -15,11 +13,11 @@ from .v2.light import async_setup_entry as setup_entry_v2
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HueConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up light entities."""
-    bridge: HueBridge = hass.data[DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
 
     if bridge.api_version == 1:
         await setup_entry_v1(hass, config_entry, async_add_entities)

homeassistant/components/hue/migration.py

@@ -10,7 +10,6 @@ from aiohue.v2.models.resource import ResourceTypes
 from homeassistant import core
 from homeassistant.components.binary_sensor import BinarySensorDeviceClass
 from homeassistant.components.sensor import SensorDeviceClass
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY, CONF_API_VERSION, CONF_HOST, CONF_USERNAME
 from homeassistant.helpers import (
     aiohttp_client,
@@ -18,12 +17,13 @@ from homeassistant.helpers import (
     entity_registry as er,
 )
 
+from .bridge import HueConfigEntry
 from .const import DOMAIN
 
 LOGGER = logging.getLogger(__name__)
 
 
-async def check_migration(hass: core.HomeAssistant, entry: ConfigEntry) -> None:
+async def check_migration(hass: core.HomeAssistant, entry: HueConfigEntry) -> None:
     """Check if config entry needs any migration actions."""
     host = entry.data[CONF_HOST]
 
@@ -66,7 +66,7 @@ async def check_migration(hass: core.HomeAssistant, entry: ConfigEntry) -> None:
             hass.config_entries.async_update_entry(entry, data=data)
 
 
-async def handle_v2_migration(hass: core.HomeAssistant, entry: ConfigEntry) -> None:
+async def handle_v2_migration(hass: core.HomeAssistant, entry: HueConfigEntry) -> None:
     """Perform migration of devices and entities to V2 Id's."""
     host = entry.data[CONF_HOST]
     api_key = entry.data[CONF_API_KEY]

homeassistant/components/hue/scene.py

@@ -12,7 +12,6 @@ from aiohue.v2.models.smart_scene import SmartScene as HueSmartScene, SmartScene
 import voluptuous as vol
 
 from homeassistant.components.scene import ATTR_TRANSITION, Scene as SceneEntity
-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 (
@@ -20,7 +19,7 @@ from homeassistant.helpers.entity_platform import (
     async_get_current_platform,
 )
 
-from .bridge import HueBridge
+from .bridge import HueBridge, HueConfigEntry
 from .const import DOMAIN
 from .v2.entity import HueBaseEntity
 from .v2.helpers import normalize_hue_brightness, normalize_hue_transition
@@ -33,11 +32,11 @@ ATTR_BRIGHTNESS = "brightness"
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HueConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up scene platform from Hue group scenes."""
-    bridge: HueBridge = hass.data[DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
     api: HueBridgeV2 = bridge.api
 
     if bridge.api_version == 1:

homeassistant/components/hue/sensor.py

@@ -2,23 +2,21 @@
 
 from __future__ import annotations
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .bridge import HueBridge
-from .const import DOMAIN
+from .bridge import HueConfigEntry
 from .v1.sensor import async_setup_entry as setup_entry_v1
 from .v2.sensor import async_setup_entry as setup_entry_v2
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HueConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up sensor entities."""
-    bridge: HueBridge = hass.data[DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
     if bridge.api_version == 1:
         await setup_entry_v1(hass, config_entry, async_add_entities)
         return

homeassistant/components/hue/services.py

@@ -12,7 +12,7 @@ from homeassistant.core import HomeAssistant, ServiceCall
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.service import verify_domain_control
 
-from .bridge import HueBridge
+from .bridge import HueBridge, HueConfigEntry
 from .const import (
     ATTR_DYNAMIC,
     ATTR_GROUP_NAME,
@@ -37,14 +37,16 @@ def async_register_services(hass: HomeAssistant) -> None:
         dynamic = call.data.get(ATTR_DYNAMIC, False)
 
         # Call the set scene function on each bridge
+        entries: list[HueConfigEntry] = hass.config_entries.async_loaded_entries(DOMAIN)
         tasks = [
-            hue_activate_scene_v1(bridge, group_name, scene_name, transition)
-            if bridge.api_version == 1
-            else hue_activate_scene_v2(
-                bridge, group_name, scene_name, transition, dynamic
+            hue_activate_scene_v1(
+                entry.runtime_data, group_name, scene_name, transition
             )
-            for bridge in hass.data[DOMAIN].values()
-            if isinstance(bridge, HueBridge)
+            if entry.runtime_data.api_version == 1
+            else hue_activate_scene_v2(
+                entry.runtime_data, group_name, scene_name, transition, dynamic
+            )
+            for entry in entries
         ]
         results = await asyncio.gather(*tasks)
 

homeassistant/components/hue/switch.py

@@ -19,23 +19,21 @@ from homeassistant.components.switch import (
     SwitchEntity,
     SwitchEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .bridge import HueBridge
-from .const import DOMAIN
+from .bridge import HueConfigEntry
 from .v2.entity import HueBaseEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HueConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Hue switch platform from Hue resources."""
-    bridge: HueBridge = hass.data[DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
     api: HueBridgeV2 = bridge.api
 
     if bridge.api_version == 1:

homeassistant/components/hue/v1/binary_sensor.py

@@ -6,16 +6,22 @@ from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
 )
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from ..const import DOMAIN as HUE_DOMAIN
+from ..bridge import HueConfigEntry
 from .sensor_base import SENSOR_CONFIG_MAP, GenericZLLSensor
 
 PRESENCE_NAME_FORMAT = "{} motion"
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: HueConfigEntry,
+    async_add_entities: AddConfigEntryEntitiesCallback,
+) -> None:
     """Defer binary sensor setup to the shared sensor module."""
-    bridge = hass.data[HUE_DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
 
     if not bridge.sensor_manager:
         return

homeassistant/components/hue/v1/device_trigger.py

@@ -27,7 +27,7 @@ from homeassistant.helpers.typing import ConfigType
 from ..const import ATTR_HUE_EVENT, CONF_SUBTYPE, DOMAIN
 
 if TYPE_CHECKING:
-    from ..bridge import HueBridge
+    from ..bridge import HueBridge, HueConfigEntry
 
 TRIGGER_SCHEMA = DEVICE_TRIGGER_BASE_SCHEMA.extend(
     {vol.Required(CONF_TYPE): str, vol.Required(CONF_SUBTYPE): str}
@@ -111,8 +111,9 @@ REMOTES: dict[str, dict[tuple[str, str], dict[str, int]]] = {
 
 def _get_hue_event_from_device_id(hass, device_id):
     """Resolve hue event from device id."""
-    for bridge in hass.data.get(DOMAIN, {}).values():
-        for hue_event in bridge.sensor_manager.current_events.values():
+    entries: list[HueConfigEntry] = hass.config_entries.async_loaded_entries(DOMAIN)
+    for entry in entries:
+        for hue_event in entry.runtime_data.sensor_manager.current_events.values():
             if device_id == hue_event.device_registry_id:
                 return hue_event
 

homeassistant/components/hue/v1/light.py

@@ -28,10 +28,11 @@ from homeassistant.components.light import (
     LightEntityFeature,
     filter_supported_color_modes,
 )
-from homeassistant.core import callback
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import PlatformNotReady
 from homeassistant.helpers.debounce import Debouncer
 from homeassistant.helpers.device_registry import DeviceInfo
+from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.update_coordinator import (
     CoordinatorEntity,
     DataUpdateCoordinator,
@@ -39,7 +40,7 @@ from homeassistant.helpers.update_coordinator import (
 )
 from homeassistant.util import color as color_util
 
-from ..bridge import HueBridge
+from ..bridge import HueConfigEntry
 from ..const import (
     CONF_ALLOW_HUE_GROUPS,
     CONF_ALLOW_UNREACHABLE,
@@ -139,11 +140,15 @@ def create_light(item_class, coordinator, bridge, is_group, rooms, api, item_id)
     )
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: HueConfigEntry,
+    async_add_entities: AddConfigEntryEntitiesCallback,
+) -> None:
     """Set up the Hue lights from a config entry."""
-    bridge: HueBridge = hass.data[DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
     api_version = tuple(int(v) for v in bridge.api.config.apiversion.split("."))
-    rooms = {}
+    rooms: dict[str, str] = {}
 
     allow_groups = config_entry.options.get(
         CONF_ALLOW_HUE_GROUPS, DEFAULT_ALLOW_HUE_GROUPS

homeassistant/components/hue/v1/sensor.py

@@ -13,8 +13,10 @@ from homeassistant.components.sensor import (
     SensorStateClass,
 )
 from homeassistant.const import LIGHT_LUX, PERCENTAGE, EntityCategory, UnitOfTemperature
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from ..const import DOMAIN as HUE_DOMAIN
+from ..bridge import HueConfigEntry
 from .sensor_base import SENSOR_CONFIG_MAP, GenericHueSensor, GenericZLLSensor
 
 LIGHT_LEVEL_NAME_FORMAT = "{} light level"
@@ -22,9 +24,13 @@ REMOTE_NAME_FORMAT = "{} battery level"
 TEMPERATURE_NAME_FORMAT = "{} temperature"
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: HueConfigEntry,
+    async_add_entities: AddConfigEntryEntitiesCallback,
+) -> None:
     """Defer sensor setup to the shared sensor module."""
-    bridge = hass.data[HUE_DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
 
     if not bridge.sensor_manager:
         return

homeassistant/components/hue/v2/binary_sensor.py

@@ -27,13 +27,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, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from ..bridge import HueBridge
-from ..const import DOMAIN
+from ..bridge import HueConfigEntry
 from .entity import HueBaseEntity
 
 type SensorType = CameraMotion | Contact | Motion | EntertainmentConfiguration | Tamper
@@ -48,11 +46,11 @@ type ControllerType = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HueConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Hue Sensors from Config Entry."""
-    bridge: HueBridge = hass.data[DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
     api: HueBridgeV2 = bridge.api
 
     @callback

homeassistant/components/hue/v2/group.py

@@ -22,14 +22,13 @@ from homeassistant.components.light import (
     LightEntityDescription,
     LightEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.util import color as color_util
 
-from ..bridge import HueBridge
+from ..bridge import HueBridge, HueConfigEntry
 from ..const import DOMAIN
 from .entity import HueBaseEntity
 from .helpers import (
@@ -41,11 +40,11 @@ from .helpers import (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HueConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Hue groups on light platform."""
-    bridge: HueBridge = hass.data[DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
     api: HueBridgeV2 = bridge.api
 
     async def async_add_light(event_type: EventType, resource: GroupedLight) -> None:

homeassistant/components/hue/v2/light.py

@@ -26,13 +26,12 @@ from homeassistant.components.light import (
     LightEntityFeature,
     filter_supported_color_modes,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
 from homeassistant.util import color as color_util
 
-from ..bridge import HueBridge
+from ..bridge import HueBridge, HueConfigEntry
 from ..const import DOMAIN
 from .entity import HueBaseEntity
 from .helpers import (
@@ -51,11 +50,11 @@ DEPRECATED_EFFECT_NONE = "None"
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HueConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Hue Light from Config Entry."""
-    bridge: HueBridge = hass.data[DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
     api: HueBridgeV2 = bridge.api
     controller: LightsController = api.lights
     make_light_entity = partial(HueLight, bridge, controller)

homeassistant/components/hue/v2/sensor.py

@@ -25,13 +25,11 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import LIGHT_LUX, PERCENTAGE, EntityCategory, UnitOfTemperature
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from ..bridge import HueBridge
-from ..const import DOMAIN
+from ..bridge import HueBridge, HueConfigEntry
 from .entity import HueBaseEntity
 
 type SensorType = DevicePower | LightLevel | Temperature | ZigbeeConnectivity
@@ -45,11 +43,11 @@ type ControllerType = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HueConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Hue Sensors from Config Entry."""
-    bridge: HueBridge = hass.data[DOMAIN][config_entry.entry_id]
+    bridge = config_entry.runtime_data
     api: HueBridgeV2 = bridge.api
     ctrl_base: SensorsController = api.sensors
 

tests/components/hue/conftest.py

@@ -59,7 +59,7 @@ def create_mock_bridge(hass: HomeAssistant, api_version: int = 1) -> Mock:
 
     async def async_initialize_bridge():
         if bridge.config_entry:
-            hass.data.setdefault(hue.DOMAIN, {})[bridge.config_entry.entry_id] = bridge
+            bridge.config_entry.runtime_data = bridge
         if bridge.api_version == 2:
             await async_setup_devices(bridge)
         return True
@@ -73,7 +73,7 @@ def create_mock_bridge(hass: HomeAssistant, api_version: int = 1) -> Mock:
 
     async def async_reset():
         if bridge.config_entry:
-            hass.data[hue.DOMAIN].pop(bridge.config_entry.entry_id)
+            delattr(bridge.config_entry, "runtime_data")
         return True
 
     bridge.async_reset = async_reset
@@ -273,7 +273,7 @@ async def setup_platform(
         api_version=mock_bridge.api_version, host=hostname
     )
     mock_bridge.config_entry = config_entry
-    hass.data[hue.DOMAIN] = {config_entry.entry_id: mock_bridge}
+    config_entry.runtime_data = {config_entry.entry_id: mock_bridge}
 
     # simulate a full setup by manually adding the bridge config entry
     await setup_bridge(hass, mock_bridge, config_entry)

tests/components/hue/test_init.py

@@ -42,7 +42,7 @@ async def test_setup_with_no_config(hass: HomeAssistant) -> None:
     assert len(hass.config_entries.flow.async_progress()) == 0
 
     # No configs stored
-    assert hue.DOMAIN not in hass.data
+    assert not hass.config_entries.async_entries(hue.DOMAIN)
 
 
 async def test_unload_entry(hass: HomeAssistant, mock_bridge_setup) -> None:
@@ -55,15 +55,15 @@ async def test_unload_entry(hass: HomeAssistant, mock_bridge_setup) -> None:
     assert await async_setup_component(hass, hue.DOMAIN, {}) is True
     assert len(mock_bridge_setup.mock_calls) == 1
 
-    hass.data[hue.DOMAIN] = {entry.entry_id: mock_bridge_setup}
+    entry.runtime_data = mock_bridge_setup
 
     async def mock_reset():
-        hass.data[hue.DOMAIN].pop(entry.entry_id)
+        delattr(entry, "runtime_data")
         return True
 
     mock_bridge_setup.async_reset = mock_reset
     assert await hue.async_unload_entry(hass, entry)
-    assert hue.DOMAIN not in hass.data
+    assert not hasattr(entry, "runtime_data")
 
 
 async def test_setting_unique_id(hass: HomeAssistant, mock_bridge_setup) -> None:

tests/components/hue/test_light_v1.py

@@ -185,7 +185,7 @@ async def setup_bridge(hass: HomeAssistant, mock_bridge_v1: Mock) -> None:
     )
     config_entry.mock_state(hass, ConfigEntryState.LOADED)
     mock_bridge_v1.config_entry = config_entry
-    hass.data[hue.DOMAIN] = {config_entry.entry_id: mock_bridge_v1}
+    config_entry.runtime_data = mock_bridge_v1
     await hass.config_entries.async_forward_entry_setups(config_entry, ["light"])
     # To flush out the service call to update the group
     await hass.async_block_till_done()