From ea0c93e3dbf03fea8a3a32cebce56c8e73bf4a46 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 22 Jun 2024 18:11:48 -0500 Subject: [PATCH] Update uiprotect to 3.1.1 (#120173) --- .../components/unifiprotect/__init__.py | 17 ++- .../components/unifiprotect/binary_sensor.py | 3 +- .../components/unifiprotect/const.py | 7 +- homeassistant/components/unifiprotect/data.py | 113 +++++++++--------- .../components/unifiprotect/manifest.json | 2 +- requirements_all.txt | 2 +- requirements_test_all.txt | 2 +- tests/components/unifiprotect/conftest.py | 8 ++ tests/components/unifiprotect/test_camera.py | 38 +++++- tests/components/unifiprotect/test_init.py | 16 ++- tests/components/unifiprotect/utils.py | 2 + 11 files changed, 123 insertions(+), 87 deletions(-) diff --git a/homeassistant/components/unifiprotect/__init__.py b/homeassistant/components/unifiprotect/__init__.py index 068c5665e6b..394a7f43329 100644 --- a/homeassistant/components/unifiprotect/__init__.py +++ b/homeassistant/components/unifiprotect/__init__.py @@ -6,6 +6,7 @@ from datetime import timedelta import logging from aiohttp.client_exceptions import ServerDisconnectedError +from uiprotect.api import DEVICE_UPDATE_INTERVAL from uiprotect.data import Bootstrap from uiprotect.data.types import FirmwareReleaseChannel from uiprotect.exceptions import ClientError, NotAuthorized @@ -29,7 +30,6 @@ from homeassistant.helpers.typing import ConfigType from .const import ( AUTH_RETRIES, CONF_ALLOW_EA, - DEFAULT_SCAN_INTERVAL, DEVICES_THAT_ADOPT, DOMAIN, MIN_REQUIRED_PROTECT_V, @@ -49,7 +49,7 @@ from .views import ThumbnailProxyView, VideoProxyView _LOGGER = logging.getLogger(__name__) -SCAN_INTERVAL = timedelta(seconds=DEFAULT_SCAN_INTERVAL) +SCAN_INTERVAL = timedelta(seconds=DEVICE_UPDATE_INTERVAL) CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN) @@ -70,11 +70,9 @@ 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") - data_service = ProtectData(hass, protect, SCAN_INTERVAL, entry) try: - bootstrap = await protect.get_bootstrap() - nvr_info = bootstrap.nvr + await protect.update() except NotAuthorized as err: retry_key = f"{entry.entry_id}_auth" retries = hass.data.setdefault(DOMAIN, {}).get(retry_key, 0) @@ -86,6 +84,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: UFPConfigEntry) -> bool: except (TimeoutError, ClientError, ServerDisconnectedError) as err: raise ConfigEntryNotReady from err + data_service = ProtectData(hass, protect, SCAN_INTERVAL, entry) + bootstrap = protect.bootstrap + nvr_info = bootstrap.nvr auth_user = bootstrap.users.get(bootstrap.auth_user_id) if auth_user and auth_user.cloud_account: ir.async_create_issue( @@ -169,11 +170,7 @@ async def _async_setup_entry( bootstrap: Bootstrap, ) -> None: await async_migrate_data(hass, entry, data_service.api, bootstrap) - - await data_service.async_setup() - if not data_service.last_update_success: - raise ConfigEntryNotReady - + data_service.async_setup() await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) hass.http.register_view(ThumbnailProxyView(hass)) hass.http.register_view(VideoProxyView(hass)) diff --git a/homeassistant/components/unifiprotect/binary_sensor.py b/homeassistant/components/unifiprotect/binary_sensor.py index 5596d3b7a62..fb60158580e 100644 --- a/homeassistant/components/unifiprotect/binary_sensor.py +++ b/homeassistant/components/unifiprotect/binary_sensor.py @@ -691,6 +691,7 @@ class ProtectDiskBinarySensor(ProtectNVREntity, BinarySensorEntity): super()._async_update_device_from_protect(device) slot = self._disk.slot self._attr_available = False + available = self.data.last_update_success # should not be possible since it would require user to # _downgrade_ to make ustorage disppear @@ -698,7 +699,7 @@ class ProtectDiskBinarySensor(ProtectNVREntity, BinarySensorEntity): for disk in self.device.system_info.ustorage.disks: if disk.slot == slot: self._disk = disk - self._attr_available = True + self._attr_available = available break self._attr_is_on = not self._disk.is_healthy diff --git a/homeassistant/components/unifiprotect/const.py b/homeassistant/components/unifiprotect/const.py index 9839d823585..b56761263f4 100644 --- a/homeassistant/components/unifiprotect/const.py +++ b/homeassistant/components/unifiprotect/const.py @@ -5,9 +5,9 @@ from uiprotect.data import ModelType, Version from homeassistant.const import Platform DOMAIN = "unifiprotect" -# some UniFi OS consoles have an unknown rate limit on auth -# if rate limit is triggered a 401 is returned -AUTH_RETRIES = 11 # ~12 hours of retries with the last waiting ~6 hours +# If rate limit for 4.x or later a 429 is returned +# so we can use a lower value +AUTH_RETRIES = 2 ATTR_EVENT_SCORE = "event_score" ATTR_EVENT_ID = "event_id" @@ -35,7 +35,6 @@ CONFIG_OPTIONS = [ DEFAULT_PORT = 443 DEFAULT_ATTRIBUTION = "Powered by UniFi Protect Server" DEFAULT_BRAND = "Ubiquiti" -DEFAULT_SCAN_INTERVAL = 60 DEFAULT_VERIFY_SSL = False DEFAULT_MAX_MEDIA = 1000 diff --git a/homeassistant/components/unifiprotect/data.py b/homeassistant/components/unifiprotect/data.py index e3e4cbc7f50..6b502eaa5f3 100644 --- a/homeassistant/components/unifiprotect/data.py +++ b/homeassistant/components/unifiprotect/data.py @@ -13,7 +13,6 @@ from typing_extensions import Generator from uiprotect import ProtectApiClient from uiprotect.data import ( NVR, - Bootstrap, Camera, Event, EventType, @@ -23,6 +22,7 @@ from uiprotect.data import ( ) from uiprotect.exceptions import ClientError, NotAuthorized from uiprotect.utils import log_event +from uiprotect.websocket import WebsocketState from homeassistant.config_entries import ConfigEntry from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback @@ -83,8 +83,7 @@ class ProtectData: str, set[Callable[[ProtectDeviceType], None]] ] = defaultdict(set) self._pending_camera_ids: set[str] = set() - self._unsub_interval: CALLBACK_TYPE | None = None - self._unsub_websocket: CALLBACK_TYPE | None = None + self._unsubs: list[CALLBACK_TYPE] = [] self._auth_failures = 0 self.last_update_success = False self.api = protect @@ -115,11 +114,9 @@ class ProtectData: self, device_types: Iterable[ModelType], ignore_unadopted: bool = True ) -> Generator[ProtectAdoptableDeviceModel]: """Get all devices matching types.""" + bootstrap = self.api.bootstrap for device_type in device_types: - devices = async_get_devices_by_type( - self.api.bootstrap, device_type - ).values() - for device in devices: + for device in async_get_devices_by_type(bootstrap, device_type).values(): if ignore_unadopted and not device.is_adopted_by_us: continue yield device @@ -130,33 +127,61 @@ class ProtectData: Generator[Camera], self.get_by_types({ModelType.CAMERA}, ignore_unadopted) ) - async def async_setup(self) -> None: + @callback + def async_setup(self) -> None: """Subscribe and do the refresh.""" - self._unsub_websocket = self.api.subscribe_websocket( - self._async_process_ws_message - ) - await self.async_refresh() + self.last_update_success = True + self._async_update_change(True, force_update=True) + api = self.api + self._unsubs = [ + api.subscribe_websocket_state(self._async_websocket_state_changed), + api.subscribe_websocket(self._async_process_ws_message), + async_track_time_interval( + self._hass, self._async_poll, self._update_interval + ), + ] + + @callback + def _async_websocket_state_changed(self, state: WebsocketState) -> None: + """Handle a change in the websocket state.""" + self._async_update_change(state is WebsocketState.CONNECTED) + + def _async_update_change( + self, + success: bool, + force_update: bool = False, + exception: Exception | None = None, + ) -> None: + """Process a change in update success.""" + was_success = self.last_update_success + self.last_update_success = success + + if not success: + level = logging.ERROR if was_success else logging.DEBUG + title = self._entry.title + _LOGGER.log(level, "%s: Connection lost", title, exc_info=exception) + self._async_process_updates() + return + + self._auth_failures = 0 + if not was_success: + _LOGGER.info("%s: Connection restored", self._entry.title) + self._async_process_updates() + elif force_update: + self._async_process_updates() async def async_stop(self, *args: Any) -> None: """Stop processing data.""" - if self._unsub_websocket: - self._unsub_websocket() - self._unsub_websocket = None - if self._unsub_interval: - self._unsub_interval() - self._unsub_interval = None + for unsub in self._unsubs: + unsub() + self._unsubs.clear() await self.api.async_disconnect_ws() - async def async_refresh(self, *_: Any, force: bool = False) -> None: + async def async_refresh(self) -> None: """Update the data.""" - - # if last update was failure, force until success - if not self.last_update_success: - force = True - try: - updates = await self.api.update(force=force) - except NotAuthorized: + await self.api.update() + except NotAuthorized as ex: if self._auth_failures < AUTH_RETRIES: _LOGGER.exception("Auth error while updating") self._auth_failures += 1 @@ -164,17 +189,11 @@ class ProtectData: await self.async_stop() _LOGGER.exception("Reauthentication required") self._entry.async_start_reauth(self._hass) - self.last_update_success = False - except ClientError: - if self.last_update_success: - _LOGGER.exception("Error while updating") - self.last_update_success = False - # manually trigger update to mark entities unavailable - self._async_process_updates(self.api.bootstrap) + self._async_update_change(False, exception=ex) + except ClientError as ex: + self._async_update_change(False, exception=ex) else: - self.last_update_success = True - self._auth_failures = 0 - self._async_process_updates(updates) + self._async_update_change(True, force_update=True) @callback def async_add_pending_camera_id(self, camera_id: str) -> None: @@ -184,7 +203,6 @@ class ProtectData: initialized yet. Will cause Websocket code to check for channels to be initialized for the camera and issue a dispatch once they do. """ - self._pending_camera_ids.add(camera_id) @callback @@ -278,25 +296,15 @@ class ProtectData: self._async_update_device(new_obj, message.changed_data) @callback - def _async_process_updates(self, updates: Bootstrap | None) -> None: + def _async_process_updates(self) -> None: """Process update from the protect data.""" - - # Websocket connected, use data from it - if updates is None: - return - self._async_signal_device_update(self.api.bootstrap.nvr) for device in self.get_by_types(DEVICES_THAT_ADOPT): self._async_signal_device_update(device) @callback def _async_poll(self, now: datetime) -> None: - """Poll the Protect API. - - If the websocket is connected, most of the time - this will be a no-op. If the websocket is disconnected, - this will trigger a reconnect and refresh. - """ + """Poll the Protect API.""" self._entry.async_create_background_task( self._hass, self.async_refresh(), @@ -309,10 +317,6 @@ class ProtectData: self, mac: str, update_callback: Callable[[ProtectDeviceType], None] ) -> CALLBACK_TYPE: """Add an callback subscriber.""" - if not self._subscriptions: - self._unsub_interval = async_track_time_interval( - self._hass, self._async_poll, self._update_interval - ) self._subscriptions[mac].add(update_callback) return partial(self._async_unsubscribe, mac, update_callback) @@ -324,9 +328,6 @@ class ProtectData: self._subscriptions[mac].remove(update_callback) if not self._subscriptions[mac]: del self._subscriptions[mac] - if not self._subscriptions and self._unsub_interval: - self._unsub_interval() - self._unsub_interval = None @callback def _async_signal_device_update(self, device: ProtectDeviceType) -> None: diff --git a/homeassistant/components/unifiprotect/manifest.json b/homeassistant/components/unifiprotect/manifest.json index 987329abbba..15b8b5b4a1b 100644 --- a/homeassistant/components/unifiprotect/manifest.json +++ b/homeassistant/components/unifiprotect/manifest.json @@ -40,7 +40,7 @@ "integration_type": "hub", "iot_class": "local_push", "loggers": ["uiprotect", "unifi_discovery"], - "requirements": ["uiprotect==1.20.0", "unifi-discovery==1.1.8"], + "requirements": ["uiprotect==3.1.1", "unifi-discovery==1.1.8"], "ssdp": [ { "manufacturer": "Ubiquiti Networks", diff --git a/requirements_all.txt b/requirements_all.txt index 0d98c4f773c..2343fa9bd4e 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -2794,7 +2794,7 @@ twitchAPI==4.0.0 uasiren==0.0.1 # homeassistant.components.unifiprotect -uiprotect==1.20.0 +uiprotect==3.1.1 # homeassistant.components.landisgyr_heat_meter ultraheat-api==0.5.7 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index ad51c91f9b6..fdd67cf9e29 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -2174,7 +2174,7 @@ twitchAPI==4.0.0 uasiren==0.0.1 # homeassistant.components.unifiprotect -uiprotect==1.20.0 +uiprotect==3.1.1 # homeassistant.components.landisgyr_heat_meter ultraheat-api==0.5.7 diff --git a/tests/components/unifiprotect/conftest.py b/tests/components/unifiprotect/conftest.py index 6366a4f9244..0bef1ff0eb9 100644 --- a/tests/components/unifiprotect/conftest.py +++ b/tests/components/unifiprotect/conftest.py @@ -29,6 +29,7 @@ from uiprotect.data import ( Viewer, WSSubscriptionMessage, ) +from uiprotect.websocket import WebsocketState from homeassistant.components.unifiprotect.const import DOMAIN from homeassistant.core import HomeAssistant @@ -148,7 +149,14 @@ def mock_entry( ufp.ws_subscription = ws_callback return Mock() + def subscribe_websocket_state( + ws_state_subscription: Callable[[WebsocketState], None], + ) -> Any: + ufp.ws_state_subscription = ws_state_subscription + return Mock() + ufp_client.subscribe_websocket = subscribe + ufp_client.subscribe_websocket_state = subscribe_websocket_state yield ufp diff --git a/tests/components/unifiprotect/test_camera.py b/tests/components/unifiprotect/test_camera.py index 444898fbd85..9fedb67fea4 100644 --- a/tests/components/unifiprotect/test_camera.py +++ b/tests/components/unifiprotect/test_camera.py @@ -4,10 +4,13 @@ from __future__ import annotations from unittest.mock import AsyncMock, Mock +from uiprotect.api import DEVICE_UPDATE_INTERVAL from uiprotect.data import Camera as ProtectCamera, CameraChannel, StateType from uiprotect.exceptions import NvrError +from uiprotect.websocket import WebsocketState from homeassistant.components.camera import ( + STATE_IDLE, CameraEntityFeature, async_get_image, async_get_stream_source, @@ -19,13 +22,13 @@ from homeassistant.components.unifiprotect.const import ( ATTR_HEIGHT, ATTR_WIDTH, DEFAULT_ATTRIBUTION, - DEFAULT_SCAN_INTERVAL, ) from homeassistant.components.unifiprotect.utils import get_camera_base_name from homeassistant.const import ( ATTR_ATTRIBUTION, ATTR_ENTITY_ID, ATTR_SUPPORTED_FEATURES, + STATE_UNAVAILABLE, Platform, ) from homeassistant.core import HomeAssistant @@ -377,7 +380,7 @@ async def test_camera_interval_update( ufp.api.bootstrap.cameras = {new_camera.id: new_camera} ufp.api.update = AsyncMock(return_value=ufp.api.bootstrap) - await time_changed(hass, DEFAULT_SCAN_INTERVAL) + await time_changed(hass, DEVICE_UPDATE_INTERVAL) state = hass.states.get(entity_id) assert state and state.state == "recording" @@ -397,19 +400,46 @@ async def test_camera_bad_interval_update( # update fails ufp.api.update = AsyncMock(side_effect=NvrError) - await time_changed(hass, DEFAULT_SCAN_INTERVAL) + await time_changed(hass, DEVICE_UPDATE_INTERVAL) state = hass.states.get(entity_id) assert state and state.state == "unavailable" # next update succeeds ufp.api.update = AsyncMock(return_value=ufp.api.bootstrap) - await time_changed(hass, DEFAULT_SCAN_INTERVAL) + await time_changed(hass, DEVICE_UPDATE_INTERVAL) state = hass.states.get(entity_id) assert state and state.state == "idle" +async def test_camera_websocket_disconnected( + hass: HomeAssistant, ufp: MockUFPFixture, camera: ProtectCamera +) -> None: + """Test the websocket gets disconnected and reconnected.""" + + await init_entry(hass, ufp, [camera]) + assert_entity_counts(hass, Platform.CAMERA, 2, 1) + entity_id = "camera.test_camera_high_resolution_channel" + + state = hass.states.get(entity_id) + assert state and state.state == STATE_IDLE + + # websocket disconnects + ufp.ws_state_subscription(WebsocketState.DISCONNECTED) + await hass.async_block_till_done() + + state = hass.states.get(entity_id) + assert state and state.state == STATE_UNAVAILABLE + + # websocket reconnects + ufp.ws_state_subscription(WebsocketState.CONNECTED) + await hass.async_block_till_done() + + state = hass.states.get(entity_id) + assert state and state.state == STATE_IDLE + + async def test_camera_ws_update( hass: HomeAssistant, ufp: MockUFPFixture, camera: ProtectCamera ) -> None: diff --git a/tests/components/unifiprotect/test_init.py b/tests/components/unifiprotect/test_init.py index 3b75afaace8..46e57c62101 100644 --- a/tests/components/unifiprotect/test_init.py +++ b/tests/components/unifiprotect/test_init.py @@ -5,12 +5,12 @@ from __future__ import annotations from unittest.mock import AsyncMock, patch from uiprotect import NotAuthorized, NvrError, ProtectApiClient +from uiprotect.api import DEVICE_UPDATE_INTERVAL from uiprotect.data import NVR, Bootstrap, CloudAccount, Light from homeassistant.components.unifiprotect.const import ( AUTH_RETRIES, CONF_DISABLE_RTSP, - DEFAULT_SCAN_INTERVAL, DOMAIN, ) from homeassistant.config_entries import ConfigEntry, ConfigEntryState @@ -116,12 +116,12 @@ async def test_setup_too_old( old_bootstrap = ufp.api.bootstrap.copy() old_bootstrap.nvr = old_nvr - ufp.api.get_bootstrap.return_value = old_bootstrap + ufp.api.update.return_value = old_bootstrap + ufp.api.bootstrap = old_bootstrap await hass.config_entries.async_setup(ufp.entry.entry_id) await hass.async_block_till_done() assert ufp.entry.state is ConfigEntryState.SETUP_ERROR - assert not ufp.api.update.called async def test_setup_cloud_account( @@ -179,13 +179,13 @@ async def test_setup_failed_update_reauth( # to verify it is not transient ufp.api.update = AsyncMock(side_effect=NotAuthorized) for _ in range(AUTH_RETRIES): - await time_changed(hass, DEFAULT_SCAN_INTERVAL) + await time_changed(hass, DEVICE_UPDATE_INTERVAL) assert len(hass.config_entries.flow._progress) == 0 assert ufp.api.update.call_count == AUTH_RETRIES assert ufp.entry.state is ConfigEntryState.LOADED - await time_changed(hass, DEFAULT_SCAN_INTERVAL) + await time_changed(hass, DEVICE_UPDATE_INTERVAL) assert ufp.api.update.call_count == AUTH_RETRIES + 1 assert len(hass.config_entries.flow._progress) == 1 @@ -193,18 +193,17 @@ async def test_setup_failed_update_reauth( async def test_setup_failed_error(hass: HomeAssistant, ufp: MockUFPFixture) -> None: """Test setup of unifiprotect entry with generic error.""" - ufp.api.get_bootstrap = AsyncMock(side_effect=NvrError) + ufp.api.update = AsyncMock(side_effect=NvrError) await hass.config_entries.async_setup(ufp.entry.entry_id) await hass.async_block_till_done() assert ufp.entry.state is ConfigEntryState.SETUP_RETRY - assert not ufp.api.update.called async def test_setup_failed_auth(hass: HomeAssistant, ufp: MockUFPFixture) -> None: """Test setup of unifiprotect entry with unauthorized error after multiple retries.""" - ufp.api.get_bootstrap = AsyncMock(side_effect=NotAuthorized) + ufp.api.update = AsyncMock(side_effect=NotAuthorized) await hass.config_entries.async_setup(ufp.entry.entry_id) assert ufp.entry.state is ConfigEntryState.SETUP_RETRY @@ -215,7 +214,6 @@ async def test_setup_failed_auth(hass: HomeAssistant, ufp: MockUFPFixture) -> No await hass.config_entries.async_reload(ufp.entry.entry_id) assert ufp.entry.state is ConfigEntryState.SETUP_ERROR - assert not ufp.api.update.called async def test_setup_starts_discovery( diff --git a/tests/components/unifiprotect/utils.py b/tests/components/unifiprotect/utils.py index ab3aefaa09d..21c01f77c5f 100644 --- a/tests/components/unifiprotect/utils.py +++ b/tests/components/unifiprotect/utils.py @@ -20,6 +20,7 @@ from uiprotect.data import ( ) from uiprotect.data.bootstrap import ProtectDeviceRef from uiprotect.test_util.anonymize import random_hex +from uiprotect.websocket import WebsocketState from homeassistant.const import Platform from homeassistant.core import HomeAssistant, split_entity_id @@ -38,6 +39,7 @@ class MockUFPFixture: entry: MockConfigEntry api: ProtectApiClient ws_subscription: Callable[[WSSubscriptionMessage], None] | None = None + ws_state_subscription: Callable[[WebsocketState], None] | None = None def ws_msg(self, msg: WSSubscriptionMessage) -> Any: """Emit WS message for testing."""