Switch to a dataclass for lutron_caseta entry data (#73500)

homeassistant/components/lutron_caseta/__init__.py

@@ -32,11 +32,8 @@ from .const import (
     ATTR_LEAP_BUTTON_NUMBER,
     ATTR_SERIAL,
     ATTR_TYPE,
-    BRIDGE_DEVICE,
     BRIDGE_DEVICE_ID,
-    BRIDGE_LEAP,
     BRIDGE_TIMEOUT,
-    BUTTON_DEVICES,
     CONF_CA_CERTS,
     CONF_CERTFILE,
     CONF_KEYFILE,
@@ -50,6 +47,7 @@ from .device_trigger import (
     DEVICE_TYPE_SUBTYPE_MAP_TO_LIP,
     LEAP_TO_DEVICE_TYPE_SUBTYPE_MAP,
 )
+from .models import LutronCasetaData
 from .util import serial_to_unique_id
 
 _LOGGER = logging.getLogger(__name__)
@@ -186,11 +184,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] = {
-        BRIDGE_LEAP: bridge,
-        BRIDGE_DEVICE: bridge_device,
-        BUTTON_DEVICES: button_devices,
-    }
+    hass.data[DOMAIN][entry_id] = LutronCasetaData(
+        bridge, bridge_device, button_devices
+    )
 
     hass.config_entries.async_setup_platforms(config_entry, PLATFORMS)
 
@@ -319,9 +315,8 @@ async def async_unload_entry(
     hass: HomeAssistant, entry: config_entries.ConfigEntry
 ) -> bool:
     """Unload the bridge bridge from a config entry."""
-    data = hass.data[DOMAIN][entry.entry_id]
-    smartbridge: Smartbridge = data[BRIDGE_LEAP]
-    await smartbridge.close()
+    data: LutronCasetaData = hass.data[DOMAIN][entry.entry_id]
+    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
@@ -402,7 +397,8 @@ async def async_remove_config_entry_device(
     hass: HomeAssistant, entry: config_entries.ConfigEntry, device_entry: dr.DeviceEntry
 ) -> bool:
     """Remove lutron_caseta config entry from a device."""
-    bridge: Smartbridge = hass.data[DOMAIN][entry.entry_id][BRIDGE_LEAP]
+    data: LutronCasetaData = hass.data[DOMAIN][entry.entry_id]
+    bridge = data.bridge
     devices = bridge.get_devices()
     buttons = bridge.buttons
     occupancy_groups = bridge.occupancy_groups

homeassistant/components/lutron_caseta/binary_sensor.py

@@ -12,7 +12,8 @@ from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import DOMAIN as CASETA_DOMAIN, LutronCasetaDevice, _area_and_name_from_name
-from .const import BRIDGE_DEVICE, BRIDGE_LEAP, CONFIG_URL, MANUFACTURER
+from .const import CONFIG_URL, MANUFACTURER
+from .models import LutronCasetaData
 
 
 async def async_setup_entry(
@@ -25,9 +26,9 @@ async def async_setup_entry(
     Adds occupancy groups from the Caseta bridge associated with the
     config_entry as binary_sensor entities.
     """
-    data = hass.data[CASETA_DOMAIN][config_entry.entry_id]
-    bridge = data[BRIDGE_LEAP]
-    bridge_device = data[BRIDGE_DEVICE]
+    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    bridge = data.bridge
+    bridge_device = data.bridge_device
     occupancy_groups = bridge.occupancy_groups
     async_add_entities(
         LutronOccupancySensor(occupancy_group, bridge, bridge_device)

homeassistant/components/lutron_caseta/const.py

@@ -10,9 +10,6 @@ STEP_IMPORT_FAILED = "import_failed"
 ERROR_CANNOT_CONNECT = "cannot_connect"
 ABORT_REASON_CANNOT_CONNECT = "cannot_connect"
 
-BRIDGE_LEAP = "leap"
-BRIDGE_DEVICE = "bridge_device"
-BUTTON_DEVICES = "button_devices"
 LUTRON_CASETA_BUTTON_EVENT = "lutron_caseta_button_event"
 
 BRIDGE_DEVICE_ID = "1"

homeassistant/components/lutron_caseta/cover.py

@@ -12,7 +12,8 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import LutronCasetaDeviceUpdatableEntity
-from .const import BRIDGE_DEVICE, BRIDGE_LEAP, DOMAIN as CASETA_DOMAIN
+from .const import DOMAIN as CASETA_DOMAIN
+from .models import LutronCasetaData
 
 
 async def async_setup_entry(
@@ -25,9 +26,9 @@ async def async_setup_entry(
     Adds shades from the Caseta bridge associated with the config_entry as
     cover entities.
     """
-    data = hass.data[CASETA_DOMAIN][config_entry.entry_id]
-    bridge = data[BRIDGE_LEAP]
-    bridge_device = data[BRIDGE_DEVICE]
+    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    bridge = data.bridge
+    bridge_device = data.bridge_device
     cover_devices = bridge.get_devices_by_domain(DOMAIN)
     async_add_entities(
         LutronCasetaCover(cover_device, bridge, bridge_device)

homeassistant/components/lutron_caseta/device_trigger.py

@@ -29,11 +29,11 @@ from .const import (
     ATTR_ACTION,
     ATTR_BUTTON_NUMBER,
     ATTR_SERIAL,
-    BUTTON_DEVICES,
     CONF_SUBTYPE,
     DOMAIN,
     LUTRON_CASETA_BUTTON_EVENT,
 )
+from .models import LutronCasetaData
 
 SUPPORTED_INPUTS_EVENTS_TYPES = [ACTION_PRESS, ACTION_RELEASE]
 
@@ -411,9 +411,9 @@ def get_button_device_by_dr_id(hass: HomeAssistant, device_id: str):
     if DOMAIN not in hass.data:
         return None
 
-    for config_entry in hass.data[DOMAIN]:
-        button_devices = hass.data[DOMAIN][config_entry][BUTTON_DEVICES]
-        if device := button_devices.get(device_id):
+    for entry_id in hass.data[DOMAIN]:
+        data: LutronCasetaData = hass.data[DOMAIN][entry_id]
+        if device := data.button_devices.get(device_id):
             return device
 
     return None

homeassistant/components/lutron_caseta/diagnostics.py

@@ -3,19 +3,19 @@ from __future__ import annotations
 
 from typing import Any
 
-from pylutron_caseta.smartbridge import Smartbridge
-
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from .const import BRIDGE_LEAP, DOMAIN
+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."""
-    bridge: Smartbridge = hass.data[DOMAIN][entry.entry_id][BRIDGE_LEAP]
+    data: LutronCasetaData = hass.data[DOMAIN][entry.entry_id]
+    bridge = data.bridge
     return {
         "entry": {
             "title": entry.title,

homeassistant/components/lutron_caseta/fan.py

@@ -13,7 +13,8 @@ from homeassistant.util.percentage import (
 )
 
 from . import LutronCasetaDeviceUpdatableEntity
-from .const import BRIDGE_DEVICE, BRIDGE_LEAP, DOMAIN as CASETA_DOMAIN
+from .const import DOMAIN as CASETA_DOMAIN
+from .models import LutronCasetaData
 
 DEFAULT_ON_PERCENTAGE = 50
 ORDERED_NAMED_FAN_SPEEDS = [FAN_LOW, FAN_MEDIUM, FAN_MEDIUM_HIGH, FAN_HIGH]
@@ -29,9 +30,9 @@ async def async_setup_entry(
     Adds fan controllers from the Caseta bridge associated with the config_entry
     as fan entities.
     """
-    data = hass.data[CASETA_DOMAIN][config_entry.entry_id]
-    bridge = data[BRIDGE_LEAP]
-    bridge_device = data[BRIDGE_DEVICE]
+    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    bridge = data.bridge
+    bridge_device = data.bridge_device
     fan_devices = bridge.get_devices_by_domain(DOMAIN)
     async_add_entities(
         LutronCasetaFan(fan_device, bridge, bridge_device) for fan_device in fan_devices

homeassistant/components/lutron_caseta/light.py

@@ -14,7 +14,8 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import LutronCasetaDeviceUpdatableEntity
-from .const import BRIDGE_DEVICE, BRIDGE_LEAP, DOMAIN as CASETA_DOMAIN
+from .const import DOMAIN as CASETA_DOMAIN
+from .models import LutronCasetaData
 
 
 def to_lutron_level(level):
@@ -37,9 +38,9 @@ async def async_setup_entry(
     Adds dimmers from the Caseta bridge associated with the config_entry as
     light entities.
     """
-    data = hass.data[CASETA_DOMAIN][config_entry.entry_id]
-    bridge = data[BRIDGE_LEAP]
-    bridge_device = data[BRIDGE_DEVICE]
+    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    bridge = data.bridge
+    bridge_device = data.bridge_device
     light_devices = bridge.get_devices_by_domain(DOMAIN)
     async_add_entities(
         LutronCasetaLight(light_device, bridge, bridge_device)

homeassistant/components/lutron_caseta/models.py

@@ -0,0 +1,18 @@
+"""The lutron_caseta integration models."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from pylutron_caseta.smartbridge import Smartbridge
+
+from homeassistant.helpers.device_registry import DeviceEntry
+
+
+@dataclass
+class LutronCasetaData:
+    """Data for the lutron_caseta integration."""
+
+    bridge: Smartbridge
+    bridge_device: dict[str, Any]
+    button_devices: dict[str, DeviceEntry]

homeassistant/components/lutron_caseta/scene.py

@@ -10,7 +10,8 @@ from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import _area_and_name_from_name
-from .const import BRIDGE_DEVICE, BRIDGE_LEAP, DOMAIN as CASETA_DOMAIN
+from .const import DOMAIN as CASETA_DOMAIN
+from .models import LutronCasetaData
 from .util import serial_to_unique_id
 
 
@@ -24,9 +25,9 @@ async def async_setup_entry(
     Adds scenes from the Caseta bridge associated with the config_entry as
     scene entities.
     """
-    data = hass.data[CASETA_DOMAIN][config_entry.entry_id]
-    bridge: Smartbridge = data[BRIDGE_LEAP]
-    bridge_device = data[BRIDGE_DEVICE]
+    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    bridge = data.bridge
+    bridge_device = data.bridge_device
     scenes = bridge.get_scenes()
     async_add_entities(
         LutronCasetaScene(scenes[scene], bridge, bridge_device) for scene in scenes

homeassistant/components/lutron_caseta/switch.py

@@ -6,7 +6,8 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import LutronCasetaDeviceUpdatableEntity
-from .const import BRIDGE_DEVICE, BRIDGE_LEAP, DOMAIN as CASETA_DOMAIN
+from .const import DOMAIN as CASETA_DOMAIN
+from .models import LutronCasetaData
 
 
 async def async_setup_entry(
@@ -19,9 +20,9 @@ async def async_setup_entry(
     Adds switches from the Caseta bridge associated with the config_entry as
     switch entities.
     """
-    data = hass.data[CASETA_DOMAIN][config_entry.entry_id]
-    bridge = data[BRIDGE_LEAP]
-    bridge_device = data[BRIDGE_DEVICE]
+    data: LutronCasetaData = hass.data[CASETA_DOMAIN][config_entry.entry_id]
+    bridge = data.bridge
+    bridge_device = data.bridge_device
     switch_devices = bridge.get_devices_by_domain(DOMAIN)
     async_add_entities(
         LutronCasetaLight(switch_device, bridge, bridge_device)

tests/components/lutron_caseta/test_device_trigger.py

@@ -1,4 +1,6 @@
 """The tests for Lutron Caséta device triggers."""
+from unittest.mock import MagicMock
+
 import pytest
 
 from homeassistant.components import automation
@@ -15,12 +17,12 @@ from homeassistant.components.lutron_caseta import (
     ATTR_TYPE,
 )
 from homeassistant.components.lutron_caseta.const import (
-    BUTTON_DEVICES,
     DOMAIN,
     LUTRON_CASETA_BUTTON_EVENT,
     MANUFACTURER,
 )
 from homeassistant.components.lutron_caseta.device_trigger import CONF_SUBTYPE
+from homeassistant.components.lutron_caseta.models import LutronCasetaData
 from homeassistant.const import CONF_DEVICE_ID, CONF_DOMAIN, CONF_PLATFORM, CONF_TYPE
 from homeassistant.helpers import device_registry
 from homeassistant.setup import async_setup_component
@@ -83,15 +85,17 @@ async def _async_setup_lutron_with_picos(hass, device_reg):
         )
         dr_button_devices[dr_device.id] = device
 
-    hass.data[DOMAIN][config_entry.entry_id] = {BUTTON_DEVICES: dr_button_devices}
-
+    hass.data[DOMAIN][config_entry.entry_id] = LutronCasetaData(
+        MagicMock(), MagicMock(), dr_button_devices
+    )
     return config_entry.entry_id
 
 
 async def test_get_triggers(hass, device_reg):
     """Test we get the expected triggers from a lutron pico."""
     config_entry_id = await _async_setup_lutron_with_picos(hass, device_reg)
-    dr_button_devices = hass.data[DOMAIN][config_entry_id][BUTTON_DEVICES]
+    data: LutronCasetaData = hass.data[DOMAIN][config_entry_id]
+    dr_button_devices = data.button_devices
     device_id = list(dr_button_devices)[0]
 
     expected_triggers = [
@@ -142,7 +146,8 @@ async def test_if_fires_on_button_event(hass, calls, device_reg):
     """Test for press trigger firing."""
 
     config_entry_id = await _async_setup_lutron_with_picos(hass, device_reg)
-    dr_button_devices = hass.data[DOMAIN][config_entry_id][BUTTON_DEVICES]
+    data: LutronCasetaData = hass.data[DOMAIN][config_entry_id]
+    dr_button_devices = data.button_devices
     device_id = list(dr_button_devices)[0]
     device = dr_button_devices[device_id]
     assert await async_setup_component(
@@ -224,7 +229,8 @@ async def test_validate_trigger_config_unknown_device(hass, calls, device_reg):
     """Test for no press with an unknown device."""
 
     config_entry_id = await _async_setup_lutron_with_picos(hass, device_reg)
-    dr_button_devices = hass.data[DOMAIN][config_entry_id][BUTTON_DEVICES]
+    data: LutronCasetaData = hass.data[DOMAIN][config_entry_id]
+    dr_button_devices = data.button_devices
     device_id = list(dr_button_devices)[0]
     device = dr_button_devices[device_id]
     device["type"] = "unknown"
@@ -267,7 +273,8 @@ async def test_validate_trigger_config_unknown_device(hass, calls, device_reg):
 async def test_validate_trigger_invalid_triggers(hass, device_reg):
     """Test for click_event with invalid triggers."""
     config_entry_id = await _async_setup_lutron_with_picos(hass, device_reg)
-    dr_button_devices = hass.data[DOMAIN][config_entry_id][BUTTON_DEVICES]
+    data: LutronCasetaData = hass.data[DOMAIN][config_entry_id]
+    dr_button_devices = data.button_devices
     device_id = list(dr_button_devices)[0]
     assert await async_setup_component(
         hass,