Refactor and unify device fetching for UniFi Protect (#77341)

homeassistant/components/unifiprotect/binary_sensor.py

@@ -10,6 +10,7 @@ from pyunifiprotect.data import (
     Camera,
     Event,
     Light,
+    ModelType,
     MountType,
     ProtectAdoptableDeviceModel,
     ProtectModelWithId,
@@ -409,12 +410,9 @@ def _async_motion_entities(
 ) -> list[ProtectDeviceEntity]:
     entities: list[ProtectDeviceEntity] = []
     devices = (
-        data.api.bootstrap.cameras.values() if ufp_device is None else [ufp_device]
+        data.get_by_types({ModelType.CAMERA}) if ufp_device is None else [ufp_device]
     )
     for device in devices:
-        if not device.is_adopted:
-            continue
-
         for description in MOTION_SENSORS:
             entities.append(ProtectEventBinarySensor(data, device, description))
             _LOGGER.debug(

homeassistant/components/unifiprotect/button.py

@@ -100,10 +100,8 @@ def _async_remove_adopt_button(
 ) -> None:
 
     entity_registry = er.async_get(hass)
-    if device.is_adopted_by_us and (
+    if entity_id := entity_registry.async_get_entity_id(
-        entity_id := entity_registry.async_get_entity_id(
         Platform.BUTTON, DOMAIN, f"{device.mac}_adopt"
-        )
     ):
         entity_registry.async_remove(entity_id)
 

homeassistant/components/unifiprotect/camera.py

@@ -3,10 +3,12 @@ from __future__ import annotations
 
 from collections.abc import Generator
 import logging
+from typing import cast
 
 from pyunifiprotect.data import (
     Camera as UFPCamera,
     CameraChannel,
+    ModelType,
     ProtectAdoptableDeviceModel,
     ProtectModelWithId,
     StateType,
@@ -42,12 +44,10 @@ def get_camera_channels(
     """Get all the camera channels."""
 
     devices = (
-        data.api.bootstrap.cameras.values() if ufp_device is None else [ufp_device]
+        data.get_by_types({ModelType.CAMERA}) if ufp_device is None else [ufp_device]
     )
     for camera in devices:
-        if not camera.is_adopted_by_us:
+        camera = cast(UFPCamera, camera)
-            continue
-
         if not camera.channels:
             if ufp_device is None:
                 # only warn on startup

homeassistant/components/unifiprotect/data.py

@@ -4,12 +4,13 @@ from __future__ import annotations
 from collections.abc import Callable, Generator, Iterable
 from datetime import timedelta
 import logging
-from typing import Any, Union
+from typing import Any, Union, cast
 
 from pyunifiprotect import ProtectApiClient
 from pyunifiprotect.data import (
     NVR,
     Bootstrap,
+    Camera,
     Event,
     EventType,
     Liveview,
@@ -35,11 +36,7 @@ from .const import (
     DISPATCH_CHANNELS,
     DOMAIN,
 )
-from .utils import (
+from .utils import async_dispatch_id as _ufpd, async_get_devices_by_type
-    async_dispatch_id as _ufpd,
-    async_get_devices,
-    async_get_devices_by_type,
-)
 
 _LOGGER = logging.getLogger(__name__)
 ProtectDeviceType = Union[ProtectAdoptableDeviceModel, NVR]
@@ -92,13 +89,17 @@ class ProtectData:
         return self._entry.options.get(CONF_MAX_MEDIA, DEFAULT_MAX_MEDIA)
 
     def get_by_types(
-        self, device_types: Iterable[ModelType]
+        self, device_types: Iterable[ModelType], ignore_unadopted: bool = True
     ) -> Generator[ProtectAdoptableDeviceModel, None, None]:
         """Get all devices matching types."""
         for device_type in device_types:
-            yield from async_get_devices_by_type(
+            devices = async_get_devices_by_type(
                 self.api.bootstrap, device_type
             ).values()
+            for device in devices:
+                if ignore_unadopted and not device.is_adopted_by_us:
+                    continue
+                yield device
 
     async def async_setup(self) -> None:
         """Subscribe and do the refresh."""
@@ -202,7 +203,8 @@ class ProtectData:
                 "Doorbell messages updated. Updating devices with LCD screens"
             )
             self.api.bootstrap.nvr.update_all_messages()
-            for camera in self.api.bootstrap.cameras.values():
+            for camera in self.get_by_types({ModelType.CAMERA}):
+                camera = cast(Camera, camera)
                 if camera.feature_flags.has_lcd_screen:
                     self._async_signal_device_update(camera)
 
@@ -250,7 +252,7 @@ class ProtectData:
             return
 
         self._async_signal_device_update(self.api.bootstrap.nvr)
-        for device in async_get_devices(self.api.bootstrap, DEVICES_THAT_ADOPT):
+        for device in self.get_by_types(DEVICES_THAT_ADOPT):
             self._async_signal_device_update(device)
 
     @callback

homeassistant/components/unifiprotect/entity.py

@@ -46,7 +46,9 @@ def _async_device_entities(
 
     entities: list[ProtectDeviceEntity] = []
     devices = (
-        [ufp_device] if ufp_device is not None else data.get_by_types({model_type})
+        [ufp_device]
+        if ufp_device is not None
+        else data.get_by_types({model_type}, ignore_unadopted=False)
     )
     for device in devices:
         assert isinstance(device, (Camera, Light, Sensor, Viewer, Doorlock, Chime))

homeassistant/components/unifiprotect/light.py

@@ -44,10 +44,7 @@ async def async_setup_entry(
     )
 
     entities = []
-    for device in data.api.bootstrap.lights.values():
+    for device in data.get_by_types({ModelType.LIGHT}):
-        if not device.is_adopted_by_us:
-            continue
-
         if device.can_write(data.api.bootstrap.auth_user):
             entities.append(ProtectLight(data, device))
 

homeassistant/components/unifiprotect/lock.py

@@ -2,11 +2,12 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
+from typing import Any, cast
 
 from pyunifiprotect.data import (
     Doorlock,
     LockStatusType,
+    ModelType,
     ProtectAdoptableDeviceModel,
     ProtectModelWithId,
 )
@@ -42,10 +43,8 @@ async def async_setup_entry(
     )
 
     entities = []
-    for device in data.api.bootstrap.doorlocks.values():
+    for device in data.get_by_types({ModelType.DOORLOCK}):
-        if not device.is_adopted_by_us:
+        device = cast(Doorlock, device)
-            continue
-
         entities.append(ProtectLock(data, device))
 
     async_add_entities(entities)

homeassistant/components/unifiprotect/media_player.py

@@ -2,9 +2,14 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
+from typing import Any, cast
 
-from pyunifiprotect.data import Camera, ProtectAdoptableDeviceModel, ProtectModelWithId
+from pyunifiprotect.data import (
+    Camera,
+    ModelType,
+    ProtectAdoptableDeviceModel,
+    ProtectModelWithId,
+)
 from pyunifiprotect.exceptions import StreamError
 
 from homeassistant.components import media_source
@@ -51,9 +56,8 @@ async def async_setup_entry(
     )
 
     entities = []
-    for device in data.api.bootstrap.cameras.values():
+    for device in data.get_by_types({ModelType.CAMERA}):
-        if not device.is_adopted_by_us:
+        device = cast(Camera, device)
-            continue
         if device.feature_flags.has_speaker:
             entities.append(ProtectMediaPlayer(data, device))
 

homeassistant/components/unifiprotect/media_source.py

@@ -7,7 +7,13 @@ from datetime import date, datetime, timedelta
 from enum import Enum
 from typing import Any, cast
 
-from pyunifiprotect.data import Camera, Event, EventType, SmartDetectObjectType
+from pyunifiprotect.data import (
+    Camera,
+    Event,
+    EventType,
+    ModelType,
+    SmartDetectObjectType,
+)
 from pyunifiprotect.exceptions import NvrError
 from pyunifiprotect.utils import from_js_time
 from yarl import URL
@@ -810,7 +816,8 @@ class ProtectMediaSource(MediaSource):
 
         cameras: list[BrowseMediaSource] = [await self._build_camera(data, "all")]
 
-        for camera in data.api.bootstrap.cameras.values():
+        for camera in data.get_by_types({ModelType.CAMERA}):
+            camera = cast(Camera, camera)
             if not camera.can_read_media(data.api.bootstrap.auth_user):
                 continue
             cameras.append(await self._build_camera(data, camera.id))

homeassistant/components/unifiprotect/sensor.py

@@ -4,13 +4,14 @@ from __future__ import annotations
 from dataclasses import dataclass
 from datetime import datetime
 import logging
-from typing import Any
+from typing import Any, cast
 
 from pyunifiprotect.data import (
     NVR,
     Camera,
     Event,
     Light,
+    ModelType,
     ProtectAdoptableDeviceModel,
     ProtectDeviceModel,
     ProtectModelWithId,
@@ -649,12 +650,10 @@ def _async_motion_entities(
 ) -> list[ProtectDeviceEntity]:
     entities: list[ProtectDeviceEntity] = []
     devices = (
-        data.api.bootstrap.cameras.values() if ufp_device is None else [ufp_device]
+        data.get_by_types({ModelType.CAMERA}) if ufp_device is None else [ufp_device]
     )
     for device in devices:
-        if not device.is_adopted_by_us:
+        device = cast(Camera, device)
-            continue
-
         for description in MOTION_TRIP_SENSORS:
             entities.append(ProtectDeviceSensor(data, device, description))
             _LOGGER.debug(