Update uiprotect to 3.1.1 (#120173)

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))

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

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
 

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:

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",

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

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

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
 
 

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:

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(

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."""