uiprotect 1.19.3__py3-none-any.whl → 2.0.0__py3-none-any.whl

uiprotect/__init__.py

@@ -4,10 +4,26 @@ from __future__ import annotations
 
 from .api import ProtectApiClient
 from .exceptions import Invalid, NotAuthorized, NvrError
+from .utils import (
+    get_nested_attr,
+    get_nested_attr_as_bool,
+    get_top_level_attr,
+    get_top_level_attr_as_bool,
+    make_enabled_getter,
+    make_required_getter,
+    make_value_getter,
+)
 
 __all__ = [
     "Invalid",
     "NotAuthorized",
     "NvrError",
     "ProtectApiClient",
+    "get_nested_attr",
+    "get_nested_attr_as_bool",
+    "get_top_level_attr",
+    "get_top_level_attr_as_bool",
+    "make_value_getter",
+    "make_enabled_getter",
+    "make_required_getter",
 ]

uiprotect/api.py

@@ -11,7 +11,7 @@ import sys
 import time
 from collections.abc import Callable
 from datetime import datetime, timedelta
-from functools import cached_property
+from functools import cached_property, partial
 from http.cookies import Morsel, SimpleCookie
 from ipaddress import IPv4Address, IPv6Address
 from pathlib import Path
@@ -203,6 +203,7 @@ class BaseApiClient:
         self._verify_ssl = verify_ssl
         self._ws_timeout = ws_timeout
         self._loaded_session = False
+        self._update_task: asyncio.Task[Bootstrap | None] | None = None
 
         self.config_dir = config_dir or (Path(user_config_dir()) / "ufp")
         self.cache_dir = cache_dir or (Path(user_cache_dir()) / "ufp_cache")
@@ -253,37 +254,56 @@ class BaseApiClient:
 
         return self._session
 
-    async def get_websocket(self) -> Websocket:
-        """Gets or creates current Websocket."""
-
-        async def _auth(force: bool) -> dict[str, str] | None:
-            if force:
-                if self._session is not None:
-                    self._session.cookie_jar.clear()
-                self.set_header("cookie", None)
-                self.set_header("x-csrf-token", None)
+    async def _auth_websocket(self, force: bool) -> dict[str, str] | None:
+        """Authenticate for Websocket."""
+        if force:
+            if self._session is not None:
+                self._session.cookie_jar.clear()
+            self.set_header("cookie", None)
+            self.set_header("x-csrf-token", None)
+            self._is_authenticated = False
 
-            await self.ensure_authenticated()
-            return self.headers
+        await self.ensure_authenticated()
+        return self.headers
 
+    def _get_websocket(self) -> Websocket:
+        """Gets or creates current Websocket."""
         if self._websocket is None:
             self._websocket = Websocket(
                 self.get_websocket_url,
-                _auth,
+                self._auth_websocket,
+                self._update_bootstrap_soon,
+                self.get_session,
+                self._process_ws_message,
                 verify=self._verify_ssl,
                 timeout=self._ws_timeout,
             )
-            self._websocket.subscribe(self._process_ws_message)
-
         return self._websocket
 
+    def _update_bootstrap_soon(self) -> None:
+        """Update bootstrap soon."""
+        _LOGGER.debug("Updating bootstrap soon")
+        # Force the next bootstrap update
+        # 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))
+
     async def close_session(self) -> None:
-        """Closing and delets client session"""
+        """Closing and deletes client session"""
+        await self._cancel_update_task()
         if self._session is not None:
             await self._session.close()
             self._session = None
             self._loaded_session = False
 
+    async def _cancel_update_task(self) -> None:
+        if self._update_task:
+            self._update_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._update_task
+            self._update_task = None
+
     def set_header(self, key: str, value: str | None) -> None:
         """Set header."""
         self.headers = self.headers or {}
@@ -413,8 +433,13 @@ class BaseApiClient:
         )
 
         if data is not None:
-            json_data: list[Any] | dict[str, Any] = orjson.loads(data)
-            return json_data
+            json_data: list[Any] | dict[str, Any]
+            try:
+                json_data = orjson.loads(data)
+                return json_data
+            except orjson.JSONDecodeError as ex:
+                _LOGGER.error("Could not decode JSON from %s", url)
+                raise NvrError(f"Could not decode JSON from {url}") from ex
         return None
 
     async def api_request_obj(
@@ -620,31 +645,17 @@ class BaseApiClient:
 
         return token_expires_at >= max_expire_time
 
-    async def async_connect_ws(self, force: bool) -> None:
-        """Connect to Websocket."""
-        if force and self._websocket is not None:
-            await self._websocket.disconnect()
-            self._websocket = None
-
-        websocket = await self.get_websocket()
-        if not websocket.is_connected:
-            self._last_ws_status = False
-            with contextlib.suppress(
-                TimeoutError,
-                asyncio.TimeoutError,
-                asyncio.CancelledError,
-            ):
-                await websocket.connect()
-
     def get_websocket_url(self) -> str:
         """Get Websocket URL."""
         return self.ws_url
 
     async def async_disconnect_ws(self) -> None:
         """Disconnect from Websocket."""
-        if self._websocket is None:
-            return
-        await self._websocket.disconnect()
+        if self._websocket:
+            websocket = self._get_websocket()
+            websocket.stop()
+            await websocket.wait_closed()
+            self._websocket = None
 
     def check_ws(self) -> bool:
         """Checks current state of Websocket."""
@@ -668,6 +679,9 @@ class BaseApiClient:
     def _get_last_update_id(self) -> str | None:
         raise NotImplementedError
 
+    async def update(self, force: bool = False) -> Bootstrap | None:
+        raise NotImplementedError
+
 
 class ProtectApiClient(BaseApiClient):
     """
@@ -748,6 +762,7 @@ class ProtectApiClient(BaseApiClient):
         self._ignore_stats = ignore_stats
         self._ws_subscriptions = []
         self.ignore_unadopted = ignore_unadopted
+        self._update_lock = asyncio.Lock()
 
         if override_connection_host:
             self._connection_host = ip_from_host(self._host)
@@ -788,41 +803,37 @@ class ProtectApiClient(BaseApiClient):
 
         You can use the various other `get_` methods if you need one off data from UFP
         """
-        now = time.monotonic()
-        now_dt = utc_now()
-        max_event_dt = now_dt - timedelta(hours=1)
-        if force:
-            self._last_update = NEVER_RAN
-            self._last_update_dt = max_event_dt
-
-        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
-            self._last_update_dt = now_dt
+        async with self._update_lock:
+            now = time.monotonic()
+            if force:
+                self._last_update = NEVER_RAN
 
-        await self.async_connect_ws(force)
-        if self.check_ws():
-            # If the websocket is connected/connecting
-            # we do not need to get events
-            _LOGGER.debug("Skipping update since websocket is active")
-            return None
+            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
+            if bootstrap_updated:
+                return None
+            self._last_update = now
+            return self._bootstrap
 
+    async def poll_events(self) -> None:
+        """Poll for events."""
+        now_dt = utc_now()
+        max_event_dt = now_dt - timedelta(hours=1)
         events = await self.get_events(
             start=self._last_update_dt or max_event_dt,
             end=now_dt,
         )
         for event in events:
             self.bootstrap.process_event(event)
-
-        self._last_update = now
         self._last_update_dt = now_dt
-        return self._bootstrap
 
     def emit_message(self, msg: WSSubscriptionMessage) -> None:
         """Emit message to all subscriptions."""
@@ -1108,13 +1119,20 @@ class ProtectApiClient(BaseApiClient):
 
         Returns a callback that will unsubscribe.
         """
-
-        def _unsub_ws_callback() -> None:
-            self._ws_subscriptions.remove(ws_callback)
-
         _LOGGER.debug("Adding subscription: %s", ws_callback)
         self._ws_subscriptions.append(ws_callback)
-        return _unsub_ws_callback
+        self._get_websocket().start()
+        return partial(self._unsubscribe_websocket, ws_callback)
+
+    def _unsubscribe_websocket(
+        self,
+        ws_callback: Callable[[WSSubscriptionMessage], None],
+    ) -> None:
+        """Unsubscribe to websocket events."""
+        _LOGGER.debug("Removing subscription: %s", ws_callback)
+        self._ws_subscriptions.remove(ws_callback)
+        if not self._ws_subscriptions:
+            self._get_websocket().stop()
 
     async def get_bootstrap(self) -> Bootstrap:
         """

uiprotect/cli/__init__.py

@@ -201,6 +201,7 @@ def shell(ctx: typer.Context) -> None:
 
     async def wait_forever() -> None:
         await protect.update()
+        protect.subscribe_websocket(lambda _: None)
         while True:
             await asyncio.sleep(10)
             await protect.update()
@@ -262,12 +263,16 @@ def profile_ws(
 
     async def callback() -> None:
         await protect.update()
+        unsub = protect.subscribe_websocket(lambda _: None)
         await profile_ws_job(
             protect,
             wait_time,
             output_path=output_path,
             ws_progress=_progress_bar,
         )
+        unsub()
+        await protect.async_disconnect_ws()
+        await protect.close_session()
 
     _setup_logger()
 

uiprotect/test_util/__init__.py

@@ -103,8 +103,10 @@ class SampleDataGenerator:
     async def async_generate(self, close_session: bool = True) -> None:
         self.log(f"Output folder: {self.output_folder}")
         self.output_folder.mkdir(parents=True, exist_ok=True)
-        websocket = await self.client.get_websocket()
-        websocket.subscribe(self._handle_ws_message)
+        websocket = self.client._get_websocket()
+        websocket.start()
+        self.log("Websocket started...")
+        websocket._subscription = self._handle_ws_message
 
         self.log("Updating devices...")
         await self.client.update()
@@ -131,8 +133,10 @@ class SampleDataGenerator:
             "chime": len(bootstrap["chimes"]),
         }
 
+        self.log("Generating event data...")
         motion_event, smart_detection = await self.generate_event_data()
         await self.generate_device_data(motion_event, smart_detection)
+        self.log("Recording websocket events...")
         await self.record_ws_events()
 
         if close_session:

uiprotect/utils.py

@@ -17,11 +17,12 @@ from copy import deepcopy
 from datetime import datetime, timedelta, timezone, tzinfo
 from decimal import Decimal
 from enum import Enum
-from functools import cache, lru_cache
+from functools import cache, lru_cache, partial
 from hashlib import sha224
 from http.cookies import Morsel
 from inspect import isclass
 from ipaddress import IPv4Address, IPv6Address, ip_address
+from operator import attrgetter
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, TypeVar, Union, overload
 from uuid import UUID
@@ -619,3 +620,57 @@ def clamp_value(value: float, step_size: float) -> float:
 def normalize_mac(mac: str) -> str:
     """Normalize MAC address."""
     return mac.lower().replace(":", "").replace("-", "").replace("_", "")
+
+
+_SENTINEL = object()
+
+
+def get_nested_attr(attrs: tuple[str, ...], obj: Any) -> Any:
+    """Fetch a nested attribute."""
+    value = obj
+    for key in attrs:
+        if (value := getattr(value, key, _SENTINEL)) is _SENTINEL:
+            return None
+    return value.value if isinstance(value, Enum) else value
+
+
+def get_nested_attr_as_bool(attrs: tuple[str, ...], obj: Any) -> bool:
+    """Fetch a nested attribute as a bool."""
+    value = obj
+    for key in attrs:
+        if (value := getattr(value, key, _SENTINEL)) is _SENTINEL:
+            return False
+    return bool(value.value if isinstance(value, Enum) else value)
+
+
+def get_top_level_attr(attr: str, obj: Any) -> Any:
+    """Fetch a top level attribute."""
+    value = getattr(obj, attr)
+    return value.value if isinstance(value, Enum) else value
+
+
+def get_top_level_attr_as_bool(attr: str, obj: Any) -> Any:
+    """Fetch a top level attribute as a bool."""
+    value = getattr(obj, attr)
+    return bool(value.value if isinstance(value, Enum) else value)
+
+
+def make_value_getter(ufp_value: str) -> Callable[[T], Any]:
+    """Return a function to get a value from a Protect device."""
+    if "." not in ufp_value:
+        return partial(get_top_level_attr, ufp_value)
+    return partial(get_nested_attr, tuple(ufp_value.split(".")))
+
+
+def make_enabled_getter(ufp_enabled: str) -> Callable[[T], bool]:
+    """Return a function to get a value from a Protect device."""
+    if "." not in ufp_enabled:
+        return attrgetter(ufp_enabled)
+    return partial(get_nested_attr, tuple(ufp_enabled.split(".")))
+
+
+def make_required_getter(ufp_required_field: str) -> Callable[[T], bool]:
+    """Return a function to get a value from a Protect device."""
+    if "." not in ufp_required_field:
+        return partial(get_top_level_attr_as_bool, ufp_required_field)
+    return partial(get_nested_attr_as_bool, tuple(ufp_required_field.split(".")))

uiprotect/websocket.py

@@ -3,9 +3,10 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import logging
-import time
-from collections.abc import Callable, Coroutine
+from collections.abc import Awaitable, Callable, Coroutine
+from http import HTTPStatus
 from typing import Any, Optional
 
 from aiohttp import (
@@ -14,213 +15,169 @@ from aiohttp import (
     ClientWebSocketResponse,
     WSMessage,
     WSMsgType,
+    WSServerHandshakeError,
 )
 
-from .utils import asyncio_timeout
-
 _LOGGER = logging.getLogger(__name__)
-CALLBACK_TYPE = Callable[..., Coroutine[Any, Any, Optional[dict[str, str]]]]
-RECENT_FAILURE_CUT_OFF = 30
-RECENT_FAILURE_THRESHOLD = 2
+AuthCallbackType = Callable[..., Coroutine[Any, Any, Optional[dict[str, str]]]]
+GetSessionCallbackType = Callable[[], Awaitable[ClientSession]]
+UpdateBootstrapCallbackType = Callable[[], None]
+_CLOSE_MESSAGE_TYPES = {WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED}
 
 
 class Websocket:
     """UniFi Protect Websocket manager."""
 
-    url: str
-    verify: bool
-    timeout_interval: int
-    backoff: int
-    _auth: CALLBACK_TYPE
-    _timeout: float
-    _ws_subscriptions: list[Callable[[WSMessage], None]]
-    _connect_lock: asyncio.Lock
-
+    _running = False
     _headers: dict[str, str] | None = None
     _websocket_loop_task: asyncio.Task[None] | None = None
-    _timer_task: asyncio.Task[None] | None = None
+    _stop_task: asyncio.Task[None] | None = None
     _ws_connection: ClientWebSocketResponse | None = None
-    _last_connect: float = -1000
-    _recent_failures: int = 0
 
     def __init__(
         self,
         get_url: Callable[[], str],
-        auth_callback: CALLBACK_TYPE,
+        auth_callback: AuthCallbackType,
+        update_bootstrap_callback: UpdateBootstrapCallbackType,
+        get_session: GetSessionCallbackType,
+        subscription: Callable[[WSMessage], None],
         *,
-        timeout: int = 30,
+        timeout: float = 30.0,
         backoff: int = 10,
         verify: bool = True,
     ) -> None:
         """Init Websocket."""
         self.get_url = get_url
-        self.timeout_interval = timeout
+        self.timeout = timeout
         self.backoff = backoff
         self.verify = verify
+        self._get_session = get_session
         self._auth = auth_callback
-        self._timeout = time.monotonic()
-        self._ws_subscriptions = []
+        self._update_bootstrap_callback = update_bootstrap_callback
         self._connect_lock = asyncio.Lock()
+        self._subscription = subscription
+        self._last_ws_connect_ok = False
 
     @property
     def is_connected(self) -> bool:
-        """Check if Websocket connected."""
-        return self._ws_connection is not None
-
-    def _get_session(self) -> ClientSession:
-        # for testing, to make easier to mock
-        return ClientSession()
+        """Return if the websocket is connected."""
+        return self._ws_connection is not None and not self._ws_connection.closed
 
-    def _process_message(self, msg: WSMessage) -> bool:
-        if msg.type == WSMsgType.ERROR:
-            _LOGGER.exception("Error from Websocket: %s", msg.data)
-            return False
+    async def _websocket_reconnect_loop(self) -> None:
+        """Reconnect loop for websocket."""
+        await self.wait_closed()
+        backoff = self.backoff
 
-        for sub in self._ws_subscriptions:
+        while True:
             try:
-                sub(msg)
+                await self._websocket_loop()
+            except ClientError:
+                _LOGGER.debug("Error in websocket reconnect loop, backoff: %s", backoff)
             except Exception:
-                _LOGGER.exception("Error processing websocket message")
+                _LOGGER.debug(
+                    "Error in websocket reconnect loop, backoff: %s",
+                    backoff,
+                    exc_info=True,
+                )
 
-        return True
+            if self._running is False:
+                break
+            await asyncio.sleep(self.backoff)
 
-    async def _websocket_loop(self, start_event: asyncio.Event) -> None:
+    async def _websocket_loop(self) -> None:
         url = self.get_url()
         _LOGGER.debug("Connecting WS to %s", url)
-        self._headers = await self._auth(self._should_reset_auth)
-
-        session = self._get_session()
+        self._headers = await self._auth(False)
+        ssl = None if self.verify else False
+        msg: WSMessage | None = None
+        seen_non_close_message = False
         # catch any and all errors for Websocket so we can clean up correctly
         try:
+            session = await self._get_session()
             self._ws_connection = await session.ws_connect(
-                url,
-                ssl=None if self.verify else False,
-                headers=self._headers,
+                url, ssl=ssl, headers=self._headers, timeout=self.timeout
             )
-            start_event.set()
-
-            self._reset_timeout()
-            async for msg in self._ws_connection:
-                if not self._process_message(msg):
+            self._last_ws_connect_ok = True
+            while True:
+                msg = await self._ws_connection.receive(self.timeout)
+                msg_type = msg.type
+                if msg_type is WSMsgType.ERROR:
+                    _LOGGER.exception("Error from Websocket: %s", msg.data)
                     break
-                self._reset_timeout()
+                elif msg_type in _CLOSE_MESSAGE_TYPES:
+                    _LOGGER.debug("Websocket closed: %s", msg)
+                    break
+
+                seen_non_close_message = True
+                try:
+                    self._subscription(msg)
+                except Exception:
+                    _LOGGER.exception("Error processing websocket message")
+        except asyncio.TimeoutError:
+            _LOGGER.debug("Websocket timeout: %s", url)
+        except WSServerHandshakeError as ex:
+            level = logging.ERROR if self._last_ws_connect_ok else logging.DEBUG
+            self._last_ws_connect_ok = False
+            if ex.status == HTTPStatus.UNAUTHORIZED.value:
+                _LOGGER.log(level, "Websocket authentication error: %s", url)
+                self._headers = await self._auth(True)
+            else:
+                _LOGGER.log(level, "Websocket handshake error: %s", url, exc_info=True)
+            raise
         except ClientError:
-            _LOGGER.exception("Websocket disconnect error: %s", url)
+            level = logging.ERROR if self._last_ws_connect_ok else logging.DEBUG
+            self._last_ws_connect_ok = False
+            _LOGGER.log(level, "Websocket disconnect error: %s", url, exc_info=True)
+            raise
         finally:
-            _LOGGER.debug("Websocket disconnected")
-            self._increase_failure()
-            self._cancel_timeout()
+            if (
+                msg is not None
+                and msg.type is WSMsgType.CLOSE
+                # If it closes right away or lastUpdateId is in the extra
+                # its an indication that we should update the bootstrap
+                # since lastUpdateId is invalid
+                and (
+                    not seen_non_close_message
+                    or (msg.extra and "lastUpdateId" in msg.extra)
+                )
+            ):
+                self._update_bootstrap_callback()
+            _LOGGER.debug("Websocket disconnected: last message: %s", msg)
             if self._ws_connection is not None and not self._ws_connection.closed:
                 await self._ws_connection.close()
-            if not session.closed:
-                await session.close()
             self._ws_connection = None
-            # make sure event does not timeout
-            start_event.set()
-
-    @property
-    def has_recent_connect(self) -> bool:
-        """Check if Websocket has recent connection."""
-        return time.monotonic() - RECENT_FAILURE_CUT_OFF <= self._last_connect
-
-    @property
-    def _should_reset_auth(self) -> bool:
-        if self.has_recent_connect:
-            if self._recent_failures > RECENT_FAILURE_THRESHOLD:
-                return True
-        else:
-            self._recent_failures = 0
-        return False
-
-    def _increase_failure(self) -> None:
-        if self.has_recent_connect:
-            self._recent_failures += 1
-        else:
-            self._recent_failures = 1
-
-    async def _do_timeout(self) -> bool:
-        _LOGGER.debug("WS timed out")
-        return await self.reconnect()
-
-    async def _timeout_loop(self) -> None:
-        while True:
-            now = time.monotonic()
-            if now > self._timeout:
-                _LOGGER.debug("WS timed out")
-                if not await self.reconnect():
-                    _LOGGER.debug("WS could not reconnect")
-                    continue
-            sleep_time = self._timeout - now
-            _LOGGER.debug("WS Timeout loop sleep %s", sleep_time)
-            await asyncio.sleep(sleep_time)
-
-    def _reset_timeout(self) -> None:
-        self._timeout = time.monotonic() + self.timeout_interval
-
-        if self._timer_task is None:
-            self._timer_task = asyncio.create_task(self._timeout_loop())
-
-    def _cancel_timeout(self) -> None:
-        if self._timer_task:
-            self._timer_task.cancel()
-
-    async def connect(self) -> bool:
-        """Connect the websocket."""
-        if self._connect_lock.locked():
-            _LOGGER.debug("Another connect is already happening")
-            return False
-        try:
-            async with asyncio_timeout(0.1):
-                await self._connect_lock.acquire()
-        except (TimeoutError, asyncio.TimeoutError, asyncio.CancelledError):
-            _LOGGER.debug("Failed to get connection lock")
 
-        start_event = asyncio.Event()
-        _LOGGER.debug("Scheduling WS connect...")
+    def start(self) -> None:
+        """Start the websocket."""
+        if self._running:
+            return
+        self._running = True
         self._websocket_loop_task = asyncio.create_task(
-            self._websocket_loop(start_event),
+            self._websocket_reconnect_loop()
         )
 
-        try:
-            async with asyncio_timeout(self.timeout_interval):
-                await start_event.wait()
-        except (TimeoutError, asyncio.TimeoutError, asyncio.CancelledError):
-            _LOGGER.warning("Timed out while waiting for Websocket to connect")
-            await self.disconnect()
-
-        self._connect_lock.release()
-        if self._ws_connection is None:
-            _LOGGER.debug("Failed to connect to Websocket")
-            return False
-        _LOGGER.debug("Connected to Websocket successfully")
-        self._last_connect = time.monotonic()
-        return True
-
-    async def disconnect(self) -> None:
+    def stop(self) -> None:
         """Disconnect the websocket."""
         _LOGGER.debug("Disconnecting websocket...")
-        if self._ws_connection is None:
+        if not self._running:
             return
-        await self._ws_connection.close()
-        self._ws_connection = None
-
-    async def reconnect(self) -> bool:
-        """Reconnect the websocket."""
-        _LOGGER.debug("Reconnecting websocket...")
-        await self.disconnect()
-        await asyncio.sleep(self.backoff)
-        return await self.connect()
-
-    def subscribe(self, ws_callback: Callable[[WSMessage], None]) -> Callable[[], None]:
-        """
-        Subscribe to raw websocket messages.
-
-        Returns a callback that will unsubscribe.
-        """
-
-        def _unsub_ws_callback() -> None:
-            self._ws_subscriptions.remove(ws_callback)
-
-        _LOGGER.debug("Adding subscription: %s", ws_callback)
-        self._ws_subscriptions.append(ws_callback)
-        return _unsub_ws_callback
+        if self._websocket_loop_task:
+            self._websocket_loop_task.cancel()
+        self._running = False
+        self._stop_task = asyncio.create_task(self._stop())
+
+    async def wait_closed(self) -> None:
+        """Wait for the websocket to close."""
+        if self._stop_task:
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._stop_task
+
+    async def _stop(self) -> None:
+        """Stop the websocket."""
+        if self._ws_connection:
+            await self._ws_connection.close()
+            self._ws_connection = None
+        if self._websocket_loop_task:
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._websocket_loop_task
+            self._websocket_loop_task = None

{uiprotect-1.19.3.dist-info → uiprotect-2.0.0.dist-info}/METADATA

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

{uiprotect-1.19.3.dist-info → uiprotect-2.0.0.dist-info}/RECORD

@@ -1,7 +1,7 @@
-uiprotect/__init__.py,sha256=llnQNtiBfwQG8IkQXovvFz4LZeFjrJx7XdmmUhu3a9E,289
+uiprotect/__init__.py,sha256=UdpRSSLSy7pdDfTKf0zRIfy6KRGt_Jv-fMzYWgibbG4,686
 uiprotect/__main__.py,sha256=C_bHCOkv5qj6WMy-6ELoY3Y6HDhLxOa1a30CzmbZhsg,462
-uiprotect/api.py,sha256=SvAni2arXGtP8HzvTljKNLclSfzpRL03RaEF8o1o6K4,65561
-uiprotect/cli/__init__.py,sha256=sSLW9keVQOkgFcMW18HTDjRrt9sJ0KWjn9DJDA6f9Pc,8658
+uiprotect/api.py,sha256=gNTK6fU3Z7PQBfs6YNTXENFxXA6RL-ayaj7XiBTDeFg,66560
+uiprotect/cli/__init__.py,sha256=1MO8rJmjjAsfVx2x01gn5DJo8B64xdPGo6gRVJbWd18,8868
 uiprotect/cli/backup.py,sha256=ZiS7RZnJGKI8TJKLW2cOUzkRM8nyTvE5Ov_jZZGtvSM,36708
 uiprotect/cli/base.py,sha256=k-_qGuNT7br0iV0KE5F4wYXF75iyLLjBEckTqxC71xM,7591
 uiprotect/cli/cameras.py,sha256=YvvMccQEYG3Wih0Ix8tan1R1vfaJ6cogg6YKWLzMUV8,16973
@@ -26,12 +26,12 @@ uiprotect/exceptions.py,sha256=kgn0cRM6lTtgLza09SDa3ZiX6ue1QqHCOogQ4qu6KTQ,965
 uiprotect/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 uiprotect/release_cache.json,sha256=NamnSFy78hOWY0DPO87J9ELFCAN6NnVquv8gQO75ZG4,386
 uiprotect/stream.py,sha256=McV3XymKyjn-1uV5jdQHcpaDjqLS4zWyMASQ8ubcyb4,4924
-uiprotect/test_util/__init__.py,sha256=d2g7afa0LSdixQ0kjEDYwafDFME_UlW2LzxpamZ2BC0,18556
+uiprotect/test_util/__init__.py,sha256=whiOUb5LfDLNT3AQG6ISiKtAqO2JnhCIdFavhWDK46M,18718
 uiprotect/test_util/anonymize.py,sha256=f-8ijU-_y9r-uAbhIPn0f0I6hzJpAkvJzc8UpWihObI,8478
-uiprotect/utils.py,sha256=NXMLqQFiTP1rHAyNK179jWgj2O0RpUCWljaLLgg8q4M,18143
-uiprotect/websocket.py,sha256=JHI_2EZeRPqPyQopsBZS0dr3tu0HaTiqeLazfBXhW_8,7339
-uiprotect-1.19.3.dist-info/LICENSE,sha256=INx18jhdbVXMEiiBANeKEbrbz57ckgzxk5uutmmcxGk,1111
-uiprotect-1.19.3.dist-info/METADATA,sha256=wwNTlD0-tf7_gRV0ORwawfuyUw5olAq8q6aadEwvXa4,10983
-uiprotect-1.19.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-uiprotect-1.19.3.dist-info/entry_points.txt,sha256=J78AUTPrTTxgI3s7SVgrmGqDP7piX2wuuEORzhDdVRA,47
-uiprotect-1.19.3.dist-info/RECORD,,
+uiprotect/utils.py,sha256=3SJFF8qs1Jz8t3mD8qwc1hFSocolFjdXI_v4yVlC7o4,20088
+uiprotect/websocket.py,sha256=enGQoS-Li9iIxcYK3KrMYtCQ4kcR1lQugiSYwROIBkU,6742
+uiprotect-2.0.0.dist-info/LICENSE,sha256=INx18jhdbVXMEiiBANeKEbrbz57ckgzxk5uutmmcxGk,1111
+uiprotect-2.0.0.dist-info/METADATA,sha256=6yuY6tDONbYf7jA3k0DDDzy2pj1LqewxXEcGGi_0mX0,10982
+uiprotect-2.0.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+uiprotect-2.0.0.dist-info/entry_points.txt,sha256=J78AUTPrTTxgI3s7SVgrmGqDP7piX2wuuEORzhDdVRA,47
+uiprotect-2.0.0.dist-info/RECORD,,