uiprotect 3.0.0__py3-none-any.whl → 3.1.0__py3-none-any.whl

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,
             )
@@ -675,22 +675,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 +684,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 +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
@@ -778,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()
 
@@ -1140,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/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

{uiprotect-3.0.0.dist-info → uiprotect-3.1.0.dist-info}/METADATA

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

{uiprotect-3.0.0.dist-info → uiprotect-3.1.0.dist-info}/RECORD

@@ -1,6 +1,6 @@
 uiprotect/__init__.py,sha256=UdpRSSLSy7pdDfTKf0zRIfy6KRGt_Jv-fMzYWgibbG4,686
 uiprotect/__main__.py,sha256=C_bHCOkv5qj6WMy-6ELoY3Y6HDhLxOa1a30CzmbZhsg,462
-uiprotect/api.py,sha256=2IV1DjkWbv4aAesGnVJkzgT5Y9DdRmIAll2FB68Jf6Q,66830
+uiprotect/api.py,sha256=JQxgQ8Tr6r3_slT0Lm0jGL5YwCE9MFjzRe3-A9aM-Qg,67560
 uiprotect/cli/__init__.py,sha256=1MO8rJmjjAsfVx2x01gn5DJo8B64xdPGo6gRVJbWd18,8868
 uiprotect/cli/backup.py,sha256=ZiS7RZnJGKI8TJKLW2cOUzkRM8nyTvE5Ov_jZZGtvSM,36708
 uiprotect/cli/base.py,sha256=k-_qGuNT7br0iV0KE5F4wYXF75iyLLjBEckTqxC71xM,7591
@@ -29,9 +29,9 @@ uiprotect/stream.py,sha256=McV3XymKyjn-1uV5jdQHcpaDjqLS4zWyMASQ8ubcyb4,4924
 uiprotect/test_util/__init__.py,sha256=whiOUb5LfDLNT3AQG6ISiKtAqO2JnhCIdFavhWDK46M,18718
 uiprotect/test_util/anonymize.py,sha256=f-8ijU-_y9r-uAbhIPn0f0I6hzJpAkvJzc8UpWihObI,8478
 uiprotect/utils.py,sha256=3SJFF8qs1Jz8t3mD8qwc1hFSocolFjdXI_v4yVlC7o4,20088
-uiprotect/websocket.py,sha256=wWcVhMST_hFamOIUQogUytMqwpsvKwKN73BYErYTrJU,6973
-uiprotect-3.0.0.dist-info/LICENSE,sha256=INx18jhdbVXMEiiBANeKEbrbz57ckgzxk5uutmmcxGk,1111
-uiprotect-3.0.0.dist-info/METADATA,sha256=slp9fl5MLh2zPWhVqhgwXczlmNS2ZtN2cGHZUDkIfJA,10982
-uiprotect-3.0.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-uiprotect-3.0.0.dist-info/entry_points.txt,sha256=J78AUTPrTTxgI3s7SVgrmGqDP7piX2wuuEORzhDdVRA,47
-uiprotect-3.0.0.dist-info/RECORD,,
+uiprotect/websocket.py,sha256=D5DZrMzo434ecp8toNxOB5HM193kVwYw42yEcg99yMw,8029
+uiprotect-3.1.0.dist-info/LICENSE,sha256=INx18jhdbVXMEiiBANeKEbrbz57ckgzxk5uutmmcxGk,1111
+uiprotect-3.1.0.dist-info/METADATA,sha256=3ocMBHOygKE3ejOejbnoVE46o-kaY5XR8dZMVEKDafk,10982
+uiprotect-3.1.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+uiprotect-3.1.0.dist-info/entry_points.txt,sha256=J78AUTPrTTxgI3s7SVgrmGqDP7piX2wuuEORzhDdVRA,47
+uiprotect-3.1.0.dist-info/RECORD,,