From 6d8f99903da97d7580d5cdf560f8064345fb9c01 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 14 Jul 2024 16:26:12 -0500
Subject: [PATCH] Migrate lutron_caseta to use entry.runtime_data (#121903)

* Migrate lutron_caseta to use entry.runtime_data

* Migrate lutron_caseta to use entry.runtime_data
---
 .../components/lutron_caseta/__init__.py      | 43 ++++++++-----------
 .../components/lutron_caseta/binary_sensor.py |  7 ++-
 .../components/lutron_caseta/button.py        |  8 ++--
 .../components/lutron_caseta/cover.py         |  8 ++--
 .../lutron_caseta/device_trigger.py           | 20 +++++----
 .../components/lutron_caseta/diagnostics.py   |  5 +--
 homeassistant/components/lutron_caseta/fan.py |  8 ++--
 .../components/lutron_caseta/light.py         |  8 +---
 .../components/lutron_caseta/models.py        |  3 ++
 .../components/lutron_caseta/scene.py         |  3 +-
 .../components/lutron_caseta/switch.py        |  4 +-
 .../lutron_caseta/test_device_trigger.py      | 18 ++++++--
 .../components/lutron_caseta/test_logbook.py  |  6 ++-
 13 files changed, 70 insertions(+), 71 deletions(-)

diff --git a/homeassistant/components/lutron_caseta/__init__.py b/homeassistant/components/lutron_caseta/__init__.py
index f6fed0688c4..178acea83f0 100644
--- a/homeassistant/components/lutron_caseta/__init__.py
+++ b/homeassistant/components/lutron_caseta/__init__.py
@@ -63,6 +63,7 @@ from .models import (
     LUTRON_KEYPAD_SERIAL,
     LUTRON_KEYPAD_TYPE,
     LutronButton,
+    LutronCasetaConfigEntry,
     LutronCasetaData,
     LutronKeypad,
     LutronKeypadData,
@@ -103,8 +104,6 @@ PLATFORMS = [
 
 async def async_setup(hass: HomeAssistant, base_config: ConfigType) -> bool:
     """Set up the Lutron component."""
-    hass.data.setdefault(DOMAIN, {})
-
     if DOMAIN in base_config:
         bridge_configs = base_config[DOMAIN]
         for config in bridge_configs:
@@ -126,7 +125,7 @@ async def async_setup(hass: HomeAssistant, base_config: ConfigType) -> bool:
 
 
 async def _async_migrate_unique_ids(
-    hass: HomeAssistant, entry: config_entries.ConfigEntry
+    hass: HomeAssistant, entry: LutronCasetaConfigEntry
 ) -> None:
     """Migrate entities since the occupancygroup were not actually unique."""
 
@@ -153,14 +152,14 @@ async def _async_migrate_unique_ids(
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, config_entry: config_entries.ConfigEntry
+    hass: HomeAssistant, entry: LutronCasetaConfigEntry
 ) -> bool:
     """Set up a bridge from a config entry."""
-    entry_id = config_entry.entry_id
-    host = config_entry.data[CONF_HOST]
-    keyfile = hass.config.path(config_entry.data[CONF_KEYFILE])
-    certfile = hass.config.path(config_entry.data[CONF_CERTFILE])
-    ca_certs = hass.config.path(config_entry.data[CONF_CA_CERTS])
+    entry_id = entry.entry_id
+    host = entry.data[CONF_HOST]
+    keyfile = hass.config.path(entry.data[CONF_KEYFILE])
+    certfile = hass.config.path(entry.data[CONF_CERTFILE])
+    ca_certs = hass.config.path(entry.data[CONF_CA_CERTS])
     bridge = None
 
     try:
@@ -185,14 +184,14 @@ async def async_setup_entry(
             raise ConfigEntryNotReady(f"Cannot connect to {host}")
 
     _LOGGER.debug("Connected to Lutron Caseta bridge via LEAP at %s", host)
-    await _async_migrate_unique_ids(hass, config_entry)
+    await _async_migrate_unique_ids(hass, entry)
 
     bridge_devices = bridge.get_devices()
     bridge_device = bridge_devices[BRIDGE_DEVICE_ID]
 
-    if not config_entry.unique_id:
+    if not entry.unique_id:
         hass.config_entries.async_update_entry(
-            config_entry, unique_id=serial_to_unique_id(bridge_device["serial"])
+            entry, unique_id=serial_to_unique_id(bridge_device["serial"])
         )
 
     _async_register_bridge_device(hass, entry_id, bridge_device, bridge)
@@ -202,13 +201,9 @@ async def async_setup_entry(
     # Store this bridge (keyed by entry_id) so it can be retrieved by the
     # platforms we're setting up.
 
-    hass.data[DOMAIN][entry_id] = LutronCasetaData(
-        bridge,
-        bridge_device,
-        keypad_data,
-    )
+    entry.runtime_data = LutronCasetaData(bridge, bridge_device, keypad_data)
 
-    await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
+    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
 
@@ -497,14 +492,12 @@ def _async_subscribe_keypad_events(
 
 
 async def async_unload_entry(
-    hass: HomeAssistant, entry: config_entries.ConfigEntry
+    hass: HomeAssistant, entry: LutronCasetaConfigEntry
 ) -> bool:
     """Unload the bridge from a config entry."""
-    data: LutronCasetaData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
     await data.bridge.close()
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
 class LutronCasetaDevice(Entity):
@@ -605,10 +598,10 @@ def _id_to_identifier(lutron_id: str) -> tuple[str, str]:
 
 
 async def async_remove_config_entry_device(
-    hass: HomeAssistant, entry: config_entries.ConfigEntry, device_entry: dr.DeviceEntry
+    hass: HomeAssistant, entry: LutronCasetaConfigEntry, device_entry: dr.DeviceEntry
 ) -> bool:
     """Remove lutron_caseta config entry from a device."""
-    data: LutronCasetaData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
     bridge = data.bridge
     devices = bridge.get_devices()
     buttons = bridge.buttons
diff --git a/homeassistant/components/lutron_caseta/binary_sensor.py b/homeassistant/components/lutron_caseta/binary_sensor.py
index 73d468a88f2..bfed8c785ae 100644
--- a/homeassistant/components/lutron_caseta/binary_sensor.py
+++ b/homeassistant/components/lutron_caseta/binary_sensor.py
@@ -6,7 +6,6 @@ from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_SUGGESTED_AREA
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
@@ -14,12 +13,12 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import DOMAIN as CASETA_DOMAIN, LutronCasetaDevice, _area_name_from_id
 from .const import CONFIG_URL, MANUFACTURER, UNASSIGNED_AREA
-from .models import LutronCasetaData
+from .models import LutronCasetaConfigEntry
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: LutronCasetaConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Lutron Caseta binary_sensor platform.
@@ -27,7 +26,7 @@ async def async_setup_entry(
     Adds occupancy groups from the Caseta bridge associated with the
     config_entry as binary_sensor entities.
     """
-    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    data = config_entry.runtime_data
     bridge = data.bridge
     occupancy_groups = bridge.occupancy_groups
     async_add_entities(
diff --git a/homeassistant/components/lutron_caseta/button.py b/homeassistant/components/lutron_caseta/button.py
index a1ed43a8b03..d2651673c4c 100644
--- a/homeassistant/components/lutron_caseta/button.py
+++ b/homeassistant/components/lutron_caseta/button.py
@@ -5,24 +5,22 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.button import ButtonEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import LutronCasetaDevice
-from .const import DOMAIN as CASETA_DOMAIN
 from .device_trigger import LEAP_TO_DEVICE_TYPE_SUBTYPE_MAP
-from .models import LutronCasetaData
+from .models import LutronCasetaConfigEntry, LutronCasetaData
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: LutronCasetaConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Lutron pico and keypad buttons."""
-    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    data = config_entry.runtime_data
     bridge = data.bridge
     button_devices = bridge.get_buttons()
     all_devices = data.bridge.get_devices()
diff --git a/homeassistant/components/lutron_caseta/cover.py b/homeassistant/components/lutron_caseta/cover.py
index 04fbb9e54c1..3edb62c0d98 100644
--- a/homeassistant/components/lutron_caseta/cover.py
+++ b/homeassistant/components/lutron_caseta/cover.py
@@ -10,13 +10,11 @@ from homeassistant.components.cover import (
     CoverEntity,
     CoverEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import LutronCasetaDeviceUpdatableEntity
-from .const import DOMAIN as CASETA_DOMAIN
-from .models import LutronCasetaData
+from .models import LutronCasetaConfigEntry
 
 
 class LutronCasetaShade(LutronCasetaDeviceUpdatableEntity, CoverEntity):
@@ -114,7 +112,7 @@ PYLUTRON_TYPE_TO_CLASSES = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: LutronCasetaConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Lutron Caseta cover platform.
@@ -122,7 +120,7 @@ async def async_setup_entry(
     Adds shades from the Caseta bridge associated with the config_entry as
     cover entities.
     """
-    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    data = config_entry.runtime_data
     bridge = data.bridge
     cover_devices = bridge.get_devices_by_domain(DOMAIN)
     async_add_entities(
diff --git a/homeassistant/components/lutron_caseta/device_trigger.py b/homeassistant/components/lutron_caseta/device_trigger.py
index 86b82e64127..0b432f88045 100644
--- a/homeassistant/components/lutron_caseta/device_trigger.py
+++ b/homeassistant/components/lutron_caseta/device_trigger.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+from typing import cast
 
 import voluptuous as vol
 
@@ -28,7 +29,7 @@ from .const import (
     DOMAIN,
     LUTRON_CASETA_BUTTON_EVENT,
 )
-from .models import LutronCasetaData
+from .models import LutronCasetaConfigEntry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -434,11 +435,14 @@ async def async_attach_trigger(
 
 def get_lutron_data_by_dr_id(hass: HomeAssistant, device_id: str):
     """Get a lutron integration data for the given device registry device id."""
-    if DOMAIN not in hass.data:
-        return None
-
-    for entry_id in hass.data[DOMAIN]:
-        data: LutronCasetaData = hass.data[DOMAIN][entry_id]
-        if data.keypad_data.dr_device_id_to_keypad.get(device_id):
-            return data
+    entries = cast(
+        list[LutronCasetaConfigEntry],
+        hass.config_entries.async_entries(
+            DOMAIN, include_ignore=False, include_disabled=False
+        ),
+    )
+    for entry in entries:
+        if hasattr(entry, "runtime_data"):
+            if entry.runtime_data.keypad_data.dr_device_id_to_keypad.get(device_id):
+                return entry.runtime_data
     return None
diff --git a/homeassistant/components/lutron_caseta/diagnostics.py b/homeassistant/components/lutron_caseta/diagnostics.py
index 61a24d21b4e..02763b14247 100644
--- a/homeassistant/components/lutron_caseta/diagnostics.py
+++ b/homeassistant/components/lutron_caseta/diagnostics.py
@@ -7,15 +7,12 @@ from typing import Any
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
-from .models import LutronCasetaData
-
 
 async def async_get_config_entry_diagnostics(
     hass: HomeAssistant, entry: ConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    data: LutronCasetaData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
     bridge = data.bridge
     return {
         "entry": {
diff --git a/homeassistant/components/lutron_caseta/fan.py b/homeassistant/components/lutron_caseta/fan.py
index 1577cf52727..cd333ba22c4 100644
--- a/homeassistant/components/lutron_caseta/fan.py
+++ b/homeassistant/components/lutron_caseta/fan.py
@@ -7,7 +7,6 @@ from typing import Any
 from pylutron_caseta import FAN_HIGH, FAN_LOW, FAN_MEDIUM, FAN_MEDIUM_HIGH, FAN_OFF
 
 from homeassistant.components.fan import DOMAIN, FanEntity, FanEntityFeature
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.percentage import (
@@ -16,8 +15,7 @@ from homeassistant.util.percentage import (
 )
 
 from . import LutronCasetaDeviceUpdatableEntity
-from .const import DOMAIN as CASETA_DOMAIN
-from .models import LutronCasetaData
+from .models import LutronCasetaConfigEntry
 
 DEFAULT_ON_PERCENTAGE = 50
 ORDERED_NAMED_FAN_SPEEDS = [FAN_LOW, FAN_MEDIUM, FAN_MEDIUM_HIGH, FAN_HIGH]
@@ -25,7 +23,7 @@ ORDERED_NAMED_FAN_SPEEDS = [FAN_LOW, FAN_MEDIUM, FAN_MEDIUM_HIGH, FAN_HIGH]
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: LutronCasetaConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Lutron Caseta fan platform.
@@ -33,7 +31,7 @@ async def async_setup_entry(
     Adds fan controllers from the Caseta bridge associated with the config_entry
     as fan entities.
     """
-    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    data = config_entry.runtime_data
     bridge = data.bridge
     fan_devices = bridge.get_devices_by_domain(DOMAIN)
     async_add_entities(LutronCasetaFan(fan_device, data) for fan_device in fan_devices)
diff --git a/homeassistant/components/lutron_caseta/light.py b/homeassistant/components/lutron_caseta/light.py
index 44c4c63e094..c0cf9449f87 100644
--- a/homeassistant/components/lutron_caseta/light.py
+++ b/homeassistant/components/lutron_caseta/light.py
@@ -25,11 +25,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import LutronCasetaDeviceUpdatableEntity
-from .const import (
-    DEVICE_TYPE_SPECTRUM_TUNE,
-    DEVICE_TYPE_WHITE_TUNE,
-    DOMAIN as CASETA_DOMAIN,
-)
+from .const import DEVICE_TYPE_SPECTRUM_TUNE, DEVICE_TYPE_WHITE_TUNE
 from .models import LutronCasetaData
 
 SUPPORTED_COLOR_MODE_DICT = {
@@ -64,7 +60,7 @@ async def async_setup_entry(
     Adds dimmers from the Caseta bridge associated with the config_entry as
     light entities.
     """
-    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    data = config_entry.runtime_data
     bridge = data.bridge
     light_devices = bridge.get_devices_by_domain(DOMAIN)
     async_add_entities(
diff --git a/homeassistant/components/lutron_caseta/models.py b/homeassistant/components/lutron_caseta/models.py
index d5ccbecbd61..402fa8885e8 100644
--- a/homeassistant/components/lutron_caseta/models.py
+++ b/homeassistant/components/lutron_caseta/models.py
@@ -8,8 +8,11 @@ from typing import Any, Final, TypedDict
 from pylutron_caseta.smartbridge import Smartbridge
 import voluptuous as vol
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.helpers.device_registry import DeviceInfo
 
+type LutronCasetaConfigEntry = ConfigEntry[LutronCasetaData]
+
 
 @dataclass
 class LutronCasetaData:
diff --git a/homeassistant/components/lutron_caseta/scene.py b/homeassistant/components/lutron_caseta/scene.py
index f4aebdafe9b..db4423495a4 100644
--- a/homeassistant/components/lutron_caseta/scene.py
+++ b/homeassistant/components/lutron_caseta/scene.py
@@ -11,7 +11,6 @@ from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN as CASETA_DOMAIN
-from .models import LutronCasetaData
 from .util import serial_to_unique_id
 
 
@@ -25,7 +24,7 @@ async def async_setup_entry(
     Adds scenes from the Caseta bridge associated with the config_entry as
     scene entities.
     """
-    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    data = config_entry.runtime_data
     bridge = data.bridge
     scenes = bridge.get_scenes()
     async_add_entities(LutronCasetaScene(scenes[scene], data) for scene in scenes)
diff --git a/homeassistant/components/lutron_caseta/switch.py b/homeassistant/components/lutron_caseta/switch.py
index 795435d5f7c..b7ec5b58b04 100644
--- a/homeassistant/components/lutron_caseta/switch.py
+++ b/homeassistant/components/lutron_caseta/switch.py
@@ -8,8 +8,6 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import LutronCasetaDeviceUpdatableEntity
-from .const import DOMAIN as CASETA_DOMAIN
-from .models import LutronCasetaData
 
 
 async def async_setup_entry(
@@ -22,7 +20,7 @@ async def async_setup_entry(
     Adds switches from the Caseta bridge associated with the config_entry as
     switch entities.
     """
-    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    data = config_entry.runtime_data
     bridge = data.bridge
     switch_devices = bridge.get_devices_by_domain(DOMAIN)
     async_add_entities(
diff --git a/tests/components/lutron_caseta/test_device_trigger.py b/tests/components/lutron_caseta/test_device_trigger.py
index 3e97be67da1..405c504dee1 100644
--- a/tests/components/lutron_caseta/test_device_trigger.py
+++ b/tests/components/lutron_caseta/test_device_trigger.py
@@ -125,7 +125,11 @@ async def _async_setup_lutron_with_picos(hass):
 async def test_get_triggers(hass: HomeAssistant) -> None:
     """Test we get the expected triggers from a lutron pico."""
     config_entry_id = await _async_setup_lutron_with_picos(hass)
-    data: LutronCasetaData = hass.data[DOMAIN][config_entry_id]
+    # Fetching the config entry runtime_data is a legacy pattern
+    # and should not be copied for new integrations
+    data: LutronCasetaData = hass.config_entries.async_get_entry(
+        config_entry_id
+    ).runtime_data
     keypads = data.keypad_data.keypads
     device_id = keypads[list(keypads)[0]]["dr_device_id"]
 
@@ -359,7 +363,11 @@ async def test_validate_trigger_config_unknown_device(
     """Test for no press with an unknown device."""
 
     config_entry_id = await _async_setup_lutron_with_picos(hass)
-    data: LutronCasetaData = hass.data[DOMAIN][config_entry_id]
+    # Fetching the config entry runtime_data is a legacy pattern
+    # and should not be copied for new integrations
+    data: LutronCasetaData = hass.config_entries.async_get_entry(
+        config_entry_id
+    ).runtime_data
     keypads = data.keypad_data.keypads
     lutron_device_id = list(keypads)[0]
     keypad = keypads[lutron_device_id]
@@ -406,7 +414,11 @@ async def test_validate_trigger_invalid_triggers(
 ) -> None:
     """Test for click_event with invalid triggers."""
     config_entry_id = await _async_setup_lutron_with_picos(hass)
-    data: LutronCasetaData = hass.data[DOMAIN][config_entry_id]
+    # Fetching the config entry runtime_data is a legacy pattern
+    # and should not be copied for new integrations
+    data: LutronCasetaData = hass.config_entries.async_get_entry(
+        config_entry_id
+    ).runtime_data
     keypads = data.keypad_data.keypads
     lutron_device_id = list(keypads)[0]
     keypad = keypads[lutron_device_id]
diff --git a/tests/components/lutron_caseta/test_logbook.py b/tests/components/lutron_caseta/test_logbook.py
index b6e8840c85c..9a58838d65c 100644
--- a/tests/components/lutron_caseta/test_logbook.py
+++ b/tests/components/lutron_caseta/test_logbook.py
@@ -53,7 +53,11 @@ async def test_humanify_lutron_caseta_button_event(hass: HomeAssistant) -> None:
 
     await hass.async_block_till_done()
 
-    data: LutronCasetaData = hass.data[DOMAIN][config_entry.entry_id]
+    # Fetching the config entry runtime_data is a legacy pattern
+    # and should not be copied for new integrations
+    data: LutronCasetaData = hass.config_entries.async_get_entry(
+        config_entry.entry_id
+    ).runtime_data
     keypads = data.keypad_data.keypads
     keypad = keypads["9"]
     dr_device_id = keypad["dr_device_id"]