Migrate unifiprotect to use entry.runtime_data (#119507)

homeassistant/components/unifiprotect/__init__.py

@@ -15,7 +15,6 @@ from uiprotect.exceptions import ClientError, NotAuthorized
 # diagnostics module will not be imported in the executor.
 from uiprotect.test_util.anonymize import anonymize_data  # noqa: F401
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
@@ -37,7 +36,7 @@ from .const import (
     OUTDATED_LOG_MESSAGE,
     PLATFORMS,
 )
-from .data import ProtectData, async_ufp_instance_for_config_entry_ids
+from .data import ProtectData, UFPConfigEntry, async_ufp_instance_for_config_entry_ids
 from .discovery import async_start_discovery
 from .migrate import async_migrate_data
 from .services import async_cleanup_services, async_setup_services
@@ -62,7 +61,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: UFPConfigEntry) -> bool:
     """Set up the UniFi Protect config entries."""
     protect = async_create_api_client(hass, entry)
     _LOGGER.debug("Connect to UniFi Protect")
@@ -107,7 +106,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     if entry.unique_id is None:
         hass.config_entries.async_update_entry(entry, unique_id=nvr_info.mac)
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = data_service
+    entry.runtime_data = data_service
     entry.async_on_unload(entry.add_update_listener(_async_options_updated))
     entry.async_on_unload(
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, data_service.async_stop)
@@ -160,7 +159,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def _async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     data_service: ProtectData,
     bootstrap: Bootstrap,
 ) -> None:
@@ -176,25 +175,24 @@ async def _async_setup_entry(
     hass.http.register_view(VideoProxyView(hass))
 
 
-async def _async_options_updated(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def _async_options_updated(hass: HomeAssistant, entry: UFPConfigEntry) -> None:
     """Update options."""
     await hass.config_entries.async_reload(entry.entry_id)
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: UFPConfigEntry) -> bool:
     """Unload UniFi Protect config entry."""
 
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        data: ProtectData = hass.data[DOMAIN][entry.entry_id]
+        data = entry.runtime_data
         await data.async_stop()
-        hass.data[DOMAIN].pop(entry.entry_id)
         async_cleanup_services(hass)
 
     return bool(unload_ok)
 
 
 async def async_remove_config_entry_device(
-    hass: HomeAssistant, config_entry: ConfigEntry, device_entry: dr.DeviceEntry
+    hass: HomeAssistant, config_entry: UFPConfigEntry, device_entry: dr.DeviceEntry
 ) -> bool:
     """Remove ufp config entry from a device."""
     unifi_macs = {

homeassistant/components/unifiprotect/binary_sensor.py

@@ -23,14 +23,13 @@ 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.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT, DOMAIN
-from .data import ProtectData
+from .const import DISPATCH_ADOPT
+from .data import ProtectData, UFPConfigEntry
 from .entity import (
     EventEntityMixin,
     ProtectDeviceEntity,
@@ -614,11 +613,11 @@ DISK_SENSORS: tuple[ProtectBinaryEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up binary sensors for UniFi Protect integration."""
-    data: ProtectData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     @callback
     def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:

homeassistant/components/unifiprotect/button.py

@@ -13,7 +13,6 @@ from homeassistant.components.button import (
     ButtonEntity,
     ButtonEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_registry as er
@@ -21,7 +20,7 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DEVICES_THAT_ADOPT, DISPATCH_ADD, DISPATCH_ADOPT, DOMAIN
-from .data import ProtectData
+from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity, async_all_device_entities
 from .models import PermRequired, ProtectSetableKeysMixin, T
 from .utils import async_dispatch_id as _ufpd
@@ -108,11 +107,11 @@ def _async_remove_adopt_button(
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Discover devices on a UniFi Protect NVR."""
-    data: ProtectData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     @callback
     def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:

homeassistant/components/unifiprotect/camera.py

@@ -16,7 +16,6 @@ from uiprotect.data import (
 )
 
 from homeassistant.components.camera import Camera, CameraEntityFeature
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import issue_registry as ir
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
@@ -33,7 +32,7 @@ from .const import (
     DISPATCH_CHANNELS,
     DOMAIN,
 )
-from .data import ProtectData
+from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity
 from .utils import async_dispatch_id as _ufpd, get_camera_base_name
 
@@ -42,7 +41,7 @@ _LOGGER = logging.getLogger(__name__)
 
 @callback
 def _create_rtsp_repair(
-    hass: HomeAssistant, entry: ConfigEntry, data: ProtectData, camera: UFPCamera
+    hass: HomeAssistant, entry: UFPConfigEntry, data: ProtectData, camera: UFPCamera
 ) -> None:
     edit_key = "readonly"
     if camera.can_write(data.api.bootstrap.auth_user):
@@ -68,7 +67,7 @@ def _create_rtsp_repair(
 @callback
 def _get_camera_channels(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     data: ProtectData,
     ufp_device: UFPCamera | None = None,
 ) -> Generator[tuple[UFPCamera, CameraChannel, bool]]:
@@ -108,7 +107,7 @@ def _get_camera_channels(
 
 def _async_camera_entities(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     data: ProtectData,
     ufp_device: UFPCamera | None = None,
 ) -> list[ProtectDeviceEntity]:
@@ -146,11 +145,11 @@ def _async_camera_entities(
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Discover cameras on a UniFi Protect NVR."""
-    data: ProtectData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     @callback
     def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:

homeassistant/components/unifiprotect/data.py

@@ -45,16 +45,15 @@ from .utils import async_dispatch_id as _ufpd, async_get_devices_by_type
 
 _LOGGER = logging.getLogger(__name__)
 type ProtectDeviceType = ProtectAdoptableDeviceModel | NVR
+type UFPConfigEntry = ConfigEntry[ProtectData]
 
 
 @callback
-def async_last_update_was_successful(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+def async_last_update_was_successful(
+    hass: HomeAssistant, entry: UFPConfigEntry
+) -> bool:
     """Check if the last update was successful for a config entry."""
-    return bool(
-        DOMAIN in hass.data
-        and entry.entry_id in hass.data[DOMAIN]
-        and hass.data[DOMAIN][entry.entry_id].last_update_success
-    )
+    return hasattr(entry, "runtime_data") and entry.runtime_data.last_update_success
 
 
 class ProtectData:
@@ -65,7 +64,7 @@ class ProtectData:
         hass: HomeAssistant,
         protect: ProtectApiClient,
         update_interval: timedelta,
-        entry: ConfigEntry,
+        entry: UFPConfigEntry,
     ) -> None:
         """Initialize an subscriber."""
         super().__init__()
@@ -316,9 +315,50 @@ def async_ufp_instance_for_config_entry_ids(
     hass: HomeAssistant, config_entry_ids: set[str]
 ) -> ProtectApiClient | None:
     """Find the UFP instance for the config entry ids."""
-    domain_data = hass.data[DOMAIN]
-    for config_entry_id in config_entry_ids:
-        if config_entry_id in domain_data:
-            protect_data: ProtectData = domain_data[config_entry_id]
-            return protect_data.api
+    return next(
+        iter(
+            entry.runtime_data.api
+            for entry_id in config_entry_ids
+            if (entry := hass.config_entries.async_get_entry(entry_id))
+        ),
+        None,
+    )
+
+
+@callback
+def async_get_ufp_entries(hass: HomeAssistant) -> list[UFPConfigEntry]:
+    """Get all the UFP entries."""
+    return cast(
+        list[UFPConfigEntry],
+        [
+            entry
+            for entry in hass.config_entries.async_entries(
+                DOMAIN, include_ignore=True, include_disabled=True
+            )
+            if hasattr(entry, "runtime_data")
+        ],
+    )
+
+
+@callback
+def async_get_data_for_nvr_id(hass: HomeAssistant, nvr_id: str) -> ProtectData | None:
+    """Find the ProtectData instance for the NVR id."""
+    return next(
+        iter(
+            entry.runtime_data
+            for entry in async_get_ufp_entries(hass)
+            if entry.runtime_data.api.bootstrap.nvr.id == nvr_id
+        ),
+        None,
+    )
+
+
+@callback
+def async_get_data_for_entry_id(
+    hass: HomeAssistant, entry_id: str
+) -> ProtectData | None:
+    """Find the ProtectData instance for a config entry id."""
+    if entry := hass.config_entries.async_get_entry(entry_id):
+        entry = cast(UFPConfigEntry, entry)
+        return entry.runtime_data
     return None

homeassistant/components/unifiprotect/diagnostics.py

@@ -6,18 +6,16 @@ from typing import Any, cast
 
 from uiprotect.test_util.anonymize import anonymize_data
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
-from .data import ProtectData
+from .data import UFPConfigEntry
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: UFPConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
 
-    data: ProtectData = hass.data[DOMAIN][config_entry.entry_id]
+    data = config_entry.runtime_data
     bootstrap = cast(dict[str, Any], anonymize_data(data.api.bootstrap.unifi_dict()))
     return {"bootstrap": bootstrap, "options": dict(config_entry.options)}

homeassistant/components/unifiprotect/light.py

@@ -13,13 +13,12 @@ from uiprotect.data import (
 )
 
 from homeassistant.components.light import ATTR_BRIGHTNESS, ColorMode, LightEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT, DOMAIN
-from .data import ProtectData
+from .const import DISPATCH_ADOPT
+from .data import UFPConfigEntry
 from .entity import ProtectDeviceEntity
 from .utils import async_dispatch_id as _ufpd
 
@@ -28,11 +27,11 @@ _LOGGER = logging.getLogger(__name__)
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up lights for UniFi Protect integration."""
-    data: ProtectData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     @callback
     def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:

homeassistant/components/unifiprotect/lock.py

@@ -14,13 +14,12 @@ from uiprotect.data import (
 )
 
 from homeassistant.components.lock import LockEntity, LockEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT, DOMAIN
-from .data import ProtectData
+from .const import DISPATCH_ADOPT
+from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity
 from .utils import async_dispatch_id as _ufpd
 
@@ -29,11 +28,11 @@ _LOGGER = logging.getLogger(__name__)
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up locks on a UniFi Protect NVR."""
-    data: ProtectData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     @callback
     def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:

homeassistant/components/unifiprotect/media_player.py

@@ -25,14 +25,13 @@ from homeassistant.components.media_player import (
     MediaType,
     async_process_play_media_url,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT, DOMAIN
-from .data import ProtectData
+from .const import DISPATCH_ADOPT
+from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity
 from .utils import async_dispatch_id as _ufpd
 
@@ -41,11 +40,11 @@ _LOGGER = logging.getLogger(__name__)
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Discover cameras with speakers on a UniFi Protect NVR."""
-    data: ProtectData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     @callback
     def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:

homeassistant/components/unifiprotect/media_source.py

@@ -26,7 +26,7 @@ from homeassistant.helpers import entity_registry as er
 from homeassistant.util import dt as dt_util
 
 from .const import DOMAIN
-from .data import ProtectData
+from .data import ProtectData, async_get_ufp_entries
 from .views import async_generate_event_video_url, async_generate_thumbnail_url
 
 VIDEO_FORMAT = "video/mp4"
@@ -89,13 +89,13 @@ def get_ufp_event(event_type: SimpleEventType) -> set[EventType]:
 
 async def async_get_media_source(hass: HomeAssistant) -> MediaSource:
     """Set up UniFi Protect media source."""
-
-    data_sources: dict[str, ProtectData] = {}
-    for data in hass.data.get(DOMAIN, {}).values():
-        if isinstance(data, ProtectData):
-            data_sources[data.api.bootstrap.nvr.id] = data
-
-    return ProtectMediaSource(hass, data_sources)
+    return ProtectMediaSource(
+        hass,
+        {
+            entry.runtime_data.api.bootstrap.nvr.id: entry.runtime_data
+            for entry in async_get_ufp_entries(hass)
+        },
+    )
 
 
 @callback

homeassistant/components/unifiprotect/migrate.py

@@ -11,13 +11,13 @@ from uiprotect.data import Bootstrap
 
 from homeassistant.components.automation import automations_with_entity
 from homeassistant.components.script import scripts_with_entity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_registry as er, issue_registry as ir
 from homeassistant.helpers.issue_registry import IssueSeverity
 
 from .const import DOMAIN
+from .data import UFPConfigEntry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -38,7 +38,7 @@ class EntityUsage(TypedDict):
 
 @callback
 def check_if_used(
-    hass: HomeAssistant, entry: ConfigEntry, entities: dict[str, EntityRef]
+    hass: HomeAssistant, entry: UFPConfigEntry, entities: dict[str, EntityRef]
 ) -> dict[str, EntityUsage]:
     """Check for usages of entities and return them."""
 
@@ -67,7 +67,7 @@ def check_if_used(
 @callback
 def create_repair_if_used(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     breaks_in: str,
     entities: dict[str, EntityRef],
 ) -> None:
@@ -101,7 +101,7 @@ def create_repair_if_used(
 
 async def async_migrate_data(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     protect: ProtectApiClient,
     bootstrap: Bootstrap,
 ) -> None:
@@ -113,7 +113,7 @@ async def async_migrate_data(
 
 
 @callback
-def async_deprecate_hdr_package(hass: HomeAssistant, entry: ConfigEntry) -> None:
+def async_deprecate_hdr_package(hass: HomeAssistant, entry: UFPConfigEntry) -> None:
     """Check for usages of hdr_mode switch and package sensor and raise repair if it is used.
 
     UniFi Protect v3.0.22 changed how HDR works so it is no longer a simple on/off toggle. There is

homeassistant/components/unifiprotect/number.py

@@ -16,14 +16,13 @@ from uiprotect.data import (
 )
 
 from homeassistant.components.number import NumberEntity, NumberEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE, EntityCategory, UnitOfTime
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT, DOMAIN
-from .data import ProtectData
+from .const import DISPATCH_ADOPT
+from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity, async_all_device_entities
 from .models import PermRequired, ProtectSetableKeysMixin, T
 from .utils import async_dispatch_id as _ufpd
@@ -220,11 +219,11 @@ CHIME_NUMBERS: tuple[ProtectNumberEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up number entities for UniFi Protect integration."""
-    data: ProtectData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     @callback
     def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:

homeassistant/components/unifiprotect/repairs.py

@@ -11,11 +11,11 @@ import voluptuous as vol
 
 from homeassistant import data_entry_flow
 from homeassistant.components.repairs import ConfirmRepairFlow, RepairsFlow
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import issue_registry as ir
 
 from .const import CONF_ALLOW_EA
+from .data import UFPConfigEntry
 from .utils import async_create_api_client
 
 
@@ -23,9 +23,9 @@ class ProtectRepair(RepairsFlow):
     """Handler for an issue fixing flow."""
 
     _api: ProtectApiClient
-    _entry: ConfigEntry
+    _entry: UFPConfigEntry
 
-    def __init__(self, *, api: ProtectApiClient, entry: ConfigEntry) -> None:
+    def __init__(self, *, api: ProtectApiClient, entry: UFPConfigEntry) -> None:
         """Create flow."""
 
         self._api = api
@@ -128,7 +128,7 @@ class RTSPRepair(ProtectRepair):
         self,
         *,
         api: ProtectApiClient,
-        entry: ConfigEntry,
+        entry: UFPConfigEntry,
         camera_id: str,
     ) -> None:
         """Create flow."""

homeassistant/components/unifiprotect/select.py

@@ -27,14 +27,13 @@ from uiprotect.data import (
 )
 
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT, DOMAIN, TYPE_EMPTY_VALUE
-from .data import ProtectData
+from .const import DISPATCH_ADOPT, TYPE_EMPTY_VALUE
+from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity, async_all_device_entities
 from .models import PermRequired, ProtectSetableKeysMixin, T
 from .utils import async_dispatch_id as _ufpd, async_get_light_motion_current
@@ -322,10 +321,10 @@ VIEWER_SELECTS: tuple[ProtectSelectEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant, entry: UFPConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up number entities for UniFi Protect integration."""
-    data: ProtectData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     @callback
     def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:

homeassistant/components/unifiprotect/sensor.py

@@ -24,7 +24,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     LIGHT_LUX,
     PERCENTAGE,
@@ -40,8 +39,8 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT, DOMAIN
-from .data import ProtectData
+from .const import DISPATCH_ADOPT
+from .data import ProtectData, UFPConfigEntry
 from .entity import (
     EventEntityMixin,
     ProtectDeviceEntity,
@@ -612,11 +611,11 @@ VIEWER_SENSORS: tuple[ProtectSensorEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up sensors for UniFi Protect integration."""
-    data: ProtectData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     @callback
     def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:

homeassistant/components/unifiprotect/switch.py

@@ -16,15 +16,14 @@ from uiprotect.data import (
 )
 
 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.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.restore_state import RestoreEntity
 
-from .const import DISPATCH_ADOPT, DOMAIN
-from .data import ProtectData
+from .const import DISPATCH_ADOPT
+from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity, ProtectNVREntity, async_all_device_entities
 from .models import PermRequired, ProtectSetableKeysMixin, T
 from .utils import async_dispatch_id as _ufpd
@@ -459,11 +458,11 @@ NVR_SWITCHES: tuple[ProtectSwitchEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up sensors for UniFi Protect integration."""
-    data: ProtectData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     @callback
     def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:

homeassistant/components/unifiprotect/text.py

@@ -13,14 +13,13 @@ from uiprotect.data import (
 )
 
 from homeassistant.components.text import TextEntity, TextEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT, DOMAIN
-from .data import ProtectData
+from .const import DISPATCH_ADOPT
+from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity, async_all_device_entities
 from .models import PermRequired, ProtectSetableKeysMixin, T
 from .utils import async_dispatch_id as _ufpd
@@ -56,11 +55,11 @@ CAMERA: tuple[ProtectTextEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: UFPConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up sensors for UniFi Protect integration."""
-    data: ProtectData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
 
     @callback
     def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:

homeassistant/components/unifiprotect/utils.py

@@ -7,7 +7,7 @@ import contextlib
 from enum import Enum
 from pathlib import Path
 import socket
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from aiohttp import CookieJar
 from typing_extensions import Generator
@@ -21,7 +21,6 @@ from uiprotect.data import (
     ProtectAdoptableDeviceModel,
 )
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -41,6 +40,9 @@ from .const import (
     ModelType,
 )
 
+if TYPE_CHECKING:
+    from .data import UFPConfigEntry
+
 _SENTINEL = object()
 
 
@@ -122,7 +124,7 @@ def async_get_light_motion_current(obj: Light) -> str:
 
 
 @callback
-def async_dispatch_id(entry: ConfigEntry, dispatch: str) -> str:
+def async_dispatch_id(entry: UFPConfigEntry, dispatch: str) -> str:
     """Generate entry specific dispatch ID."""
 
     return f"{DOMAIN}.{entry.entry_id}.{dispatch}"
@@ -130,7 +132,7 @@ def async_dispatch_id(entry: ConfigEntry, dispatch: str) -> str:
 
 @callback
 def async_create_api_client(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: UFPConfigEntry
 ) -> ProtectApiClient:
     """Create ProtectApiClient from config entry."""
 

homeassistant/components/unifiprotect/views.py

@@ -16,8 +16,7 @@ from homeassistant.components.http import HomeAssistantView
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 
-from .const import DOMAIN
-from .data import ProtectData
+from .data import ProtectData, async_get_data_for_entry_id, async_get_data_for_nvr_id
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -99,18 +98,13 @@ class ProtectProxyView(HomeAssistantView):
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize a thumbnail proxy view."""
         self.hass = hass
-        self.data = hass.data[DOMAIN]
 
-    def _get_data_or_404(self, nvr_id: str) -> ProtectData | web.Response:
-        all_data: list[ProtectData] = []
-
-        for entry_id, data in self.data.items():
-            if isinstance(data, ProtectData):
-                if nvr_id == entry_id:
-                    return data
-                if data.api.bootstrap.nvr.id == nvr_id:
-                    return data
-                all_data.append(data)
+    def _get_data_or_404(self, nvr_id_or_entry_id: str) -> ProtectData | web.Response:
+        if data := (
+            async_get_data_for_nvr_id(self.hass, nvr_id_or_entry_id)
+            or async_get_data_for_entry_id(self.hass, nvr_id_or_entry_id)
+        ):
+            return data
         return _404("Invalid NVR ID")
 
 

tests/components/unifiprotect/test_services.py

@@ -17,6 +17,7 @@ from homeassistant.components.unifiprotect.services import (
     SERVICE_SET_CHIME_PAIRED,
     SERVICE_SET_DEFAULT_DOORBELL_TEXT,
 )
+from homeassistant.config_entries import ConfigEntryDisabler
 from homeassistant.const import ATTR_DEVICE_ID, ATTR_ENTITY_ID, ATTR_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
@@ -142,6 +143,29 @@ async def test_set_default_doorbell_text(
     nvr.set_default_doorbell_message.assert_called_once_with("Test Message")
 
 
+async def test_add_doorbell_text_disabled_config_entry(
+    hass: HomeAssistant, device: dr.DeviceEntry, ufp: MockUFPFixture
+) -> None:
+    """Test add_doorbell_text service."""
+    nvr = ufp.api.bootstrap.nvr
+    nvr.__fields__["add_custom_doorbell_message"] = Mock(final=False)
+    nvr.add_custom_doorbell_message = AsyncMock()
+
+    await hass.config_entries.async_set_disabled_by(
+        ufp.entry.entry_id, ConfigEntryDisabler.USER
+    )
+    await hass.async_block_till_done()
+
+    with pytest.raises(HomeAssistantError):
+        await hass.services.async_call(
+            DOMAIN,
+            SERVICE_ADD_DOORBELL_TEXT,
+            {ATTR_DEVICE_ID: device.id, ATTR_MESSAGE: "Test Message"},
+            blocking=True,
+        )
+    assert not nvr.add_custom_doorbell_message.called
+
+
 async def test_set_chime_paired_doorbells(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,

tests/components/unifiprotect/test_views.py

@@ -149,6 +149,25 @@ async def test_thumbnail_entry_id(
     ufp.api.get_event_thumbnail.assert_called_with("test_id", width=None, height=None)
 
 
+async def test_thumbnail_invalid_entry_entry_id(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    ufp: MockUFPFixture,
+    camera: Camera,
+) -> None:
+    """Test invalid config entry ID in URL."""
+
+    ufp.api.get_event_thumbnail = AsyncMock(return_value=b"testtest")
+
+    await init_entry(hass, ufp, [camera])
+    url = async_generate_thumbnail_url("test_id", "invalid")
+
+    http_client = await hass_client()
+    response = cast(ClientResponse, await http_client.get(url))
+
+    assert response.status == 404
+
+
 async def test_video_bad_event(
     hass: HomeAssistant,
     ufp: MockUFPFixture,