diff --git a/homeassistant/components/unifiprotect/__init__.py b/homeassistant/components/unifiprotect/__init__.py index 0f41011361d..38e45798789 100644 --- a/homeassistant/components/unifiprotect/__init__.py +++ b/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 = { diff --git a/homeassistant/components/unifiprotect/binary_sensor.py b/homeassistant/components/unifiprotect/binary_sensor.py index 7e66f5efb28..c97197fea5e 100644 --- a/homeassistant/components/unifiprotect/binary_sensor.py +++ b/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: diff --git a/homeassistant/components/unifiprotect/button.py b/homeassistant/components/unifiprotect/button.py index 0db05a6cdc9..009f9b275dc 100644 --- a/homeassistant/components/unifiprotect/button.py +++ b/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: diff --git a/homeassistant/components/unifiprotect/camera.py b/homeassistant/components/unifiprotect/camera.py index 04ac2a823a3..5a703dc5458 100644 --- a/homeassistant/components/unifiprotect/camera.py +++ b/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: diff --git a/homeassistant/components/unifiprotect/data.py b/homeassistant/components/unifiprotect/data.py index 5ca9b5aaeb7..4e63ff01bc7 100644 --- a/homeassistant/components/unifiprotect/data.py +++ b/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 diff --git a/homeassistant/components/unifiprotect/diagnostics.py b/homeassistant/components/unifiprotect/diagnostics.py index ac651f6138d..b72f35db0b5 100644 --- a/homeassistant/components/unifiprotect/diagnostics.py +++ b/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)} diff --git a/homeassistant/components/unifiprotect/light.py b/homeassistant/components/unifiprotect/light.py index 18e611f2307..e119a4a59d5 100644 --- a/homeassistant/components/unifiprotect/light.py +++ b/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: diff --git a/homeassistant/components/unifiprotect/lock.py b/homeassistant/components/unifiprotect/lock.py index 6bb1dd7b4ee..4deeafa0782 100644 --- a/homeassistant/components/unifiprotect/lock.py +++ b/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: diff --git a/homeassistant/components/unifiprotect/media_player.py b/homeassistant/components/unifiprotect/media_player.py index eb17137842b..f3761b5c18a 100644 --- a/homeassistant/components/unifiprotect/media_player.py +++ b/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: diff --git a/homeassistant/components/unifiprotect/media_source.py b/homeassistant/components/unifiprotect/media_source.py index 1a67efcfd03..9d94c3ecda7 100644 --- a/homeassistant/components/unifiprotect/media_source.py +++ b/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 diff --git a/homeassistant/components/unifiprotect/migrate.py b/homeassistant/components/unifiprotect/migrate.py index a95341f497a..e469b684518 100644 --- a/homeassistant/components/unifiprotect/migrate.py +++ b/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 diff --git a/homeassistant/components/unifiprotect/number.py b/homeassistant/components/unifiprotect/number.py index ceb8614e77e..2a8137f50f7 100644 --- a/homeassistant/components/unifiprotect/number.py +++ b/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: diff --git a/homeassistant/components/unifiprotect/repairs.py b/homeassistant/components/unifiprotect/repairs.py index 3cc8967ea0d..0e505f87391 100644 --- a/homeassistant/components/unifiprotect/repairs.py +++ b/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.""" diff --git a/homeassistant/components/unifiprotect/select.py b/homeassistant/components/unifiprotect/select.py index f4a9d58e346..5ba557a8af6 100644 --- a/homeassistant/components/unifiprotect/select.py +++ b/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: diff --git a/homeassistant/components/unifiprotect/sensor.py b/homeassistant/components/unifiprotect/sensor.py index 00849c095f0..a69e9d48293 100644 --- a/homeassistant/components/unifiprotect/sensor.py +++ b/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: diff --git a/homeassistant/components/unifiprotect/switch.py b/homeassistant/components/unifiprotect/switch.py index 50953e2b8fe..d13c49af882 100644 --- a/homeassistant/components/unifiprotect/switch.py +++ b/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: diff --git a/homeassistant/components/unifiprotect/text.py b/homeassistant/components/unifiprotect/text.py index 05e6712fa65..c267419bd6d 100644 --- a/homeassistant/components/unifiprotect/text.py +++ b/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: diff --git a/homeassistant/components/unifiprotect/utils.py b/homeassistant/components/unifiprotect/utils.py index 5a0809ef9ac..ad4c99379c8 100644 --- a/homeassistant/components/unifiprotect/utils.py +++ b/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.""" diff --git a/homeassistant/components/unifiprotect/views.py b/homeassistant/components/unifiprotect/views.py index b359fd5d948..00128492c67 100644 --- a/homeassistant/components/unifiprotect/views.py +++ b/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") diff --git a/tests/components/unifiprotect/test_services.py b/tests/components/unifiprotect/test_services.py index 0a90a2d5667..b468c2de9a8 100644 --- a/tests/components/unifiprotect/test_services.py +++ b/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, diff --git a/tests/components/unifiprotect/test_views.py b/tests/components/unifiprotect/test_views.py index 6d190eb4dd6..2b80a41b16f 100644 --- a/tests/components/unifiprotect/test_views.py +++ b/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,