uiprotect 2.3.0__tar.gz → 3.1.0__tar.gz

{uiprotect-2.3.0 → uiprotect-3.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: uiprotect
-Version: 2.3.0
+Version: 3.1.0
 Summary: Python API for Unifi Protect (Unofficial)
 Home-page: https://github.com/uilibs/uiprotect
 License: MIT

{uiprotect-2.3.0 → uiprotect-3.1.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "uiprotect"
-version = "2.3.0"
+version = "3.1.0"
 description = "Python API for Unifi Protect (Unofficial)"
 authors = ["UI Protect Maintainers <ui@koston.org>"]
 license = "MIT"

{uiprotect-2.3.0 → uiprotect-3.1.0}/src/uiprotect/api.py

@@ -64,7 +64,7 @@ from .utils import (
     to_js_time,
     utc_now,
 )
-from .websocket import Websocket
+from .websocket import Websocket, WebsocketState
 
 if sys.version_info[:2] < (3, 13):
     from http import cookies
@@ -75,7 +75,6 @@ if sys.version_info[:2] < (3, 13):
 
 TOKEN_COOKIE_MAX_EXP_SECONDS = 60
 
-NEVER_RAN = -1000
 # how many seconds before the bootstrap is refreshed from Protect
 DEVICE_UPDATE_INTERVAL = 900
 # retry timeout for thumbnails/heatmaps
@@ -164,8 +163,6 @@ class BaseApiClient:
     _ws_timeout: int
 
     _is_authenticated: bool = False
-    _last_update: float = NEVER_RAN
-    _last_ws_status: bool = False
     _last_token_cookie: Morsel[str] | None = None
     _last_token_cookie_decode: dict[str, Any] | None = None
     _session: aiohttp.ClientSession | None = None
@@ -277,6 +274,7 @@ class BaseApiClient:
                 self._update_bootstrap_soon,
                 self.get_session,
                 self._process_ws_message,
+                self._on_websocket_state_change,
                 verify=self._verify_ssl,
                 timeout=self._ws_timeout,
             )
@@ -289,7 +287,7 @@ class BaseApiClient:
         # since the lastUpdateId is not valid anymore
         if self._update_task and not self._update_task.done():
             return
-        self._update_task = asyncio.create_task(self.update(force=True))
+        self._update_task = asyncio.create_task(self.update())
 
     async def close_session(self) -> None:
         """Closing and deletes client session"""
@@ -677,31 +675,19 @@ class BaseApiClient:
             await websocket.wait_closed()
             self._websocket = None
 
-    def check_ws(self) -> bool:
-        """Checks current state of Websocket."""
-        if self._websocket is None:
-            return False
-
-        if not self._websocket.is_connected:
-            log = _LOGGER.debug
-            if self._last_ws_status:
-                log = _LOGGER.warning
-            log("Websocket connection not active, failing back to polling")
-        elif not self._last_ws_status:
-            _LOGGER.info("Websocket re-connected successfully")
-
-        self._last_ws_status = self._websocket.is_connected
-        return self._last_ws_status
-
     def _process_ws_message(self, msg: aiohttp.WSMessage) -> None:
         raise NotImplementedError
 
     def _get_last_update_id(self) -> str | None:
         raise NotImplementedError
 
-    async def update(self, force: bool = False) -> Bootstrap | None:
+    async def update(self) -> Bootstrap:
         raise NotImplementedError
 
+    def _on_websocket_state_change(self, state: WebsocketState) -> None:
+        """Websocket state changed."""
+        _LOGGER.debug("Websocket state changed: %s", state)
+
 
 class ProtectApiClient(BaseApiClient):
     """
@@ -710,8 +696,7 @@ class ProtectApiClient(BaseApiClient):
     UniFi Protect is a full async application. "normal" use of interacting with it is
     to call `.update()` which will initialize the `.bootstrap` and create a Websocket
     connection to UFP. This Websocket connection will emit messages that will automatically
-    update the `.bootstrap` over time. Caling `.udpate` again (without `force`) will
-    verify the integry of the Websocket connection.
+    update the `.bootstrap` over time.
 
     You can use the `.get_` methods to one off pull devices from the UFP API, but should
     not be used for building an aplication on top of.
@@ -739,6 +724,7 @@ class ProtectApiClient(BaseApiClient):
     _subscribed_models: set[ModelType]
     _ignore_stats: bool
     _ws_subscriptions: list[Callable[[WSSubscriptionMessage], None]]
+    _ws_state_subscriptions: list[Callable[[WebsocketState], None]]
     _bootstrap: Bootstrap | None = None
     _last_update_dt: datetime | None = None
     _connection_host: IPv4Address | IPv6Address | str | None = None
@@ -781,6 +767,7 @@ class ProtectApiClient(BaseApiClient):
         self._subscribed_models = subscribed_models or set()
         self._ignore_stats = ignore_stats
         self._ws_subscriptions = []
+        self._ws_state_subscriptions = []
         self.ignore_unadopted = ignore_unadopted
         self._update_lock = asyncio.Lock()
 
@@ -816,32 +803,21 @@ class ProtectApiClient(BaseApiClient):
 
         return self._connection_host
 
-    async def update(self, force: bool = False) -> Bootstrap | None:
+    async def update(self) -> Bootstrap:
         """
-        Updates the state of devices, initalizes `.bootstrap` and
-        connects to UFP Websocket for real time updates
+        Updates the state of devices, initializes `.bootstrap`
+
+        The websocket is auto connected once there are any
+        subscriptions to it. update must be called at least
+        once before subscribing to the websocket.
 
         You can use the various other `get_` methods if you need one off data from UFP
         """
         async with self._update_lock:
-            now = time.monotonic()
-            if force:
-                self._last_update = NEVER_RAN
-
-            bootstrap_updated = False
-            if (
-                self._bootstrap is None
-                or now - self._last_update > DEVICE_UPDATE_INTERVAL
-            ):
-                bootstrap_updated = True
-                self._bootstrap = await self.get_bootstrap()
-                self.__dict__.pop("bootstrap", None)
-                self._last_update = now
-
-            if bootstrap_updated:
-                return None
-            self._last_update = now
-            return self._bootstrap
+            bootstrap = await self.get_bootstrap()
+            self.__dict__.pop("bootstrap", None)
+            self._bootstrap = bootstrap
+            return bootstrap
 
     async def poll_events(self) -> None:
         """Poll for events."""
@@ -1154,6 +1130,34 @@ class ProtectApiClient(BaseApiClient):
         if not self._ws_subscriptions:
             self._get_websocket().stop()
 
+    def subscribe_websocket_state(
+        self,
+        ws_callback: Callable[[WebsocketState], None],
+    ) -> Callable[[], None]:
+        """
+        Subscribe to websocket state changes.
+
+        Returns a callback that will unsubscribe.
+        """
+        self._ws_state_subscriptions.append(ws_callback)
+        return partial(self._unsubscribe_websocket_state, ws_callback)
+
+    def _unsubscribe_websocket_state(
+        self,
+        ws_callback: Callable[[WebsocketState], None],
+    ) -> None:
+        """Unsubscribe to websocket state changes."""
+        self._ws_state_subscriptions.remove(ws_callback)
+
+    def _on_websocket_state_change(self, state: WebsocketState) -> None:
+        """Websocket state changed."""
+        super()._on_websocket_state_change(state)
+        for sub in self._ws_state_subscriptions:
+            try:
+                sub(state)
+            except Exception:
+                _LOGGER.exception("Exception while running websocket state handler")
+
     async def get_bootstrap(self) -> Bootstrap:
         """
         Gets bootstrap object from UFP instance

{uiprotect-2.3.0 → uiprotect-3.1.0}/src/uiprotect/websocket.py

@@ -6,6 +6,7 @@ import asyncio
 import contextlib
 import logging
 from collections.abc import Awaitable, Callable, Coroutine
+from enum import Enum
 from http import HTTPStatus
 from typing import Any, Optional
 
@@ -28,6 +29,11 @@ UpdateBootstrapCallbackType = Callable[[], None]
 _CLOSE_MESSAGE_TYPES = {WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED}
 
 
+class WebsocketState(Enum):
+    CONNECTED = True
+    DISCONNECTED = False
+
+
 class Websocket:
     """UniFi Protect Websocket manager."""
 
@@ -44,6 +50,7 @@ class Websocket:
         update_bootstrap: UpdateBootstrapCallbackType,
         get_session: GetSessionCallbackType,
         subscription: Callable[[WSMessage], None],
+        state_callback: Callable[[WebsocketState], None],
         *,
         timeout: float = 30.0,
         backoff: int = 10,
@@ -59,10 +66,12 @@ class Websocket:
         self._update_bootstrap = update_bootstrap
         self._subscription = subscription
         self._seen_non_close_message = False
+        self._websocket_state = state_callback
+        self._current_state: WebsocketState = WebsocketState.DISCONNECTED
 
     @property
     def is_connected(self) -> bool:
-        """Return if the websocket is connected."""
+        """Return if the websocket is connected and has received a valid message."""
         return self._ws_connection is not None and not self._ws_connection.closed
 
     async def _websocket_loop(self) -> None:
@@ -92,11 +101,19 @@ class Websocket:
             except Exception:
                 _LOGGER.exception("Unexpected error in websocket loop")
 
+            self._state_changed(WebsocketState.DISCONNECTED)
             if self._running is False:
                 break
             _LOGGER.debug("Reconnecting websocket in %s seconds", backoff)
             await asyncio.sleep(self.backoff)
 
+    def _state_changed(self, state: WebsocketState) -> None:
+        """State changed."""
+        if self._current_state is state:
+            return
+        self._current_state = state
+        self._websocket_state(state)
+
     async def _websocket_inner_loop(self, url: URL) -> None:
         _LOGGER.debug("Connecting WS to %s", url)
         await self._attempt_auth(False)
@@ -119,7 +136,9 @@ class Websocket:
                     _LOGGER.debug("Websocket closed: %s", msg)
                     break
 
-                self._seen_non_close_message = True
+                if not self._seen_non_close_message:
+                    self._seen_non_close_message = True
+                    self._state_changed(WebsocketState.CONNECTED)
                 try:
                     self._subscription(msg)
                 except Exception:
@@ -166,21 +185,30 @@ class Websocket:
         if self._websocket_loop_task:
             self._websocket_loop_task.cancel()
         self._running = False
-        self._stop_task = asyncio.create_task(self._stop())
+        ws_connection = self._ws_connection
+        websocket_loop_task = self._websocket_loop_task
+        self._ws_connection = None
+        self._websocket_loop_task = None
+        self._stop_task = asyncio.create_task(
+            self._stop(ws_connection, websocket_loop_task)
+        )
+        self._state_changed(WebsocketState.DISCONNECTED)
 
     async def wait_closed(self) -> None:
         """Wait for the websocket to close."""
-        if self._stop_task:
+        if self._stop_task and not self._stop_task.done():
             with contextlib.suppress(asyncio.CancelledError):
                 await self._stop_task
             self._stop_task = None
 
-    async def _stop(self) -> None:
+    async def _stop(
+        self,
+        ws_connection: ClientWebSocketResponse | None,
+        websocket_loop_task: asyncio.Task[None] | None,
+    ) -> None:
         """Stop the websocket."""
-        if self._ws_connection:
-            await self._ws_connection.close()
-            self._ws_connection = None
-        if self._websocket_loop_task:
+        if ws_connection:
+            await ws_connection.close()
+        if websocket_loop_task:
             with contextlib.suppress(asyncio.CancelledError):
-                await self._websocket_loop_task
-            self._websocket_loop_task = None
+                await websocket_loop_task