uiprotect 3.0.0__tar.gz → 3.1.1__tar.gz

{uiprotect-3.0.0 → uiprotect-3.1.1}/PKG-INFO

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

{uiprotect-3.0.0 → uiprotect-3.1.1}/pyproject.toml

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

{uiprotect-3.0.0 → uiprotect-3.1.1}/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
@@ -163,7 +163,6 @@ class BaseApiClient:
     _ws_timeout: int
 
     _is_authenticated: bool = False
-    _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
@@ -275,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,
             )
@@ -411,11 +411,13 @@ class BaseApiClient:
     async def _raise_for_status(
         self, response: aiohttp.ClientResponse, raise_exception: bool = True
     ) -> None:
+        """Raise an exception based on the response status."""
         url = response.url
         reason = await get_response_reason(response)
         msg = "Request failed: %s - Status: %s - Reason: %s"
+        status = response.status
+
         if raise_exception:
-            status = response.status
             if status in {
                 HTTPStatus.UNAUTHORIZED.value,
                 HTTPStatus.FORBIDDEN.value,
@@ -430,6 +432,7 @@ class BaseApiClient:
             ):
                 raise BadRequest(msg % (url, status, reason))
             raise NvrError(msg % (url, status, reason))
+
         _LOGGER.debug(msg, url, status, reason)
 
     async def api_request(
@@ -675,22 +678,6 @@ 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
 
@@ -700,6 +687,10 @@ class BaseApiClient:
     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):
     """
@@ -736,6 +727,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
@@ -778,6 +770,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()
 
@@ -1140,6 +1133,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-3.0.0 → uiprotect-3.1.1}/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