uiprotect 1.20.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

uiprotect/api.py

@@ -11,7 +11,8 @@ 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 import HTTPStatus
 from http.cookies import Morsel, SimpleCookie
 from ipaddress import IPv4Address, IPv6Address
 from pathlib import Path
@@ -203,6 +204,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")
@@ -221,23 +223,24 @@ class BaseApiClient:
         """Updates the url after changing _host or _port."""
         if self._port != 443:
             self._url = URL(f"https://{self._host}:{self._port}")
+            self._ws_url = URL(f"wss://{self._host}:{self._port}{self.ws_path}")
         else:
             self._url = URL(f"https://{self._host}")
+            self._ws_url = URL(f"wss://{self._host}{self.ws_path}")
 
         self.base_url = str(self._url)
 
     @property
-    def ws_url(self) -> str:
+    def _ws_url_object(self) -> URL:
         """Get Websocket URL."""
-        url = f"wss://{self._host}"
-        if self._port != 443:
-            url += f":{self._port}"
+        if last_update_id := self._get_last_update_id():
+            return self._ws_url.with_query(lastUpdateId=last_update_id)
+        return self._ws_url
 
-        url += self.ws_path
-        last_update_id = self._get_last_update_id()
-        if last_update_id is None:
-            return url
-        return f"{url}?lastUpdateId={last_update_id}"
+    @property
+    def ws_url(self) -> str:
+        """Get Websocket URL."""
+        return str(self._ws_url_object)
 
     @property
     def config_file(self) -> Path:
@@ -253,37 +256,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._get_websocket_url,
+                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 {}
@@ -375,15 +397,7 @@ class BaseApiClient:
 
         try:
             if response.status != 200:
-                reason = await get_response_reason(response)
-                msg = "Request failed: %s - Status: %s - Reason: %s"
-                if raise_exception:
-                    if response.status in {401, 403}:
-                        raise NotAuthorized(msg % (url, response.status, reason))
-                    if response.status >= 400 and response.status < 500:
-                        raise BadRequest(msg % (url, response.status, reason))
-                    raise NvrError(msg % (url, response.status, reason))
-                _LOGGER.debug(msg, url, response.status, reason)
+                await self._raise_for_status(response, raise_exception)
                 return None
 
             data: bytes | None = await response.read()
@@ -396,6 +410,30 @@ class BaseApiClient:
             # re-raise exception
             raise
 
+    async def _raise_for_status(
+        self, response: aiohttp.ClientResponse, raise_exception: bool = True
+    ) -> None:
+        url = response.url
+        reason = await get_response_reason(response)
+        msg = "Request failed: %s - Status: %s - Reason: %s"
+        if raise_exception:
+            status = response.status
+            if status in {
+                HTTPStatus.UNAUTHORIZED.value,
+                HTTPStatus.FORBIDDEN.value,
+            }:
+                raise NotAuthorized(msg % (url, status, reason))
+            elif status == HTTPStatus.TOO_MANY_REQUESTS.value:
+                _LOGGER.debug("Too many requests - Login is rate limited: %s", response)
+                raise NvrError(msg % (url, status, reason))
+            elif (
+                status >= HTTPStatus.BAD_REQUEST.value
+                and status < HTTPStatus.INTERNAL_SERVER_ERROR.value
+            ):
+                raise BadRequest(msg % (url, status, reason))
+            raise NvrError(msg % (url, status, reason))
+        _LOGGER.debug(msg, url, status, reason)
+
     async def api_request(
         self,
         url: str,
@@ -413,8 +451,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(
@@ -487,6 +530,8 @@ class BaseApiClient:
             }
 
             response = await self.request("post", url=url, json=auth)
+            if response.status != 200:
+                await self._raise_for_status(response, True)
             self.set_header("cookie", response.headers.get("set-cookie", ""))
             self._is_authenticated = True
             _LOGGER.debug("Authenticated successfully!")
@@ -620,31 +665,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:
+    def _get_websocket_url(self) -> URL:
         """Get Websocket URL."""
-        return self.ws_url
+        return self._ws_url_object
 
     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 +699,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 +782,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 +823,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 +1139,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/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,173 @@ from aiohttp import (
     ClientWebSocketResponse,
     WSMessage,
     WSMsgType,
+    WSServerHandshakeError,
 )
-
-from .utils import asyncio_timeout
+from yarl import URL
 
 _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,
+        get_url: Callable[[], URL],
+        auth_callback: AuthCallbackType,
+        update_bootstrap: 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._connect_lock = asyncio.Lock()
+        self._update_bootstrap = update_bootstrap
+        self._subscription = subscription
+        self._seen_non_close_message = 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_loop(self) -> None:
+        """Running loop for websocket."""
+        await self.wait_closed()
+        backoff = self.backoff
 
-        for sub in self._ws_subscriptions:
+        while True:
+            url = self.get_url()
             try:
-                sub(msg)
+                await self._websocket_inner_loop(url)
+            except ClientError as ex:
+                level = logging.ERROR if self._seen_non_close_message else logging.DEBUG
+                if isinstance(ex, WSServerHandshakeError):
+                    if ex.status == HTTPStatus.UNAUTHORIZED.value:
+                        _LOGGER.log(level, "Websocket authentication error: %s", url)
+                        await self._attempt_reauth()
+                    else:
+                        _LOGGER.log(
+                            level, "Websocket handshake error: %s", url, exc_info=True
+                        )
+                else:
+                    _LOGGER.log(
+                        level, "Websocket disconnect error: %s", url, exc_info=True
+                    )
+            except asyncio.TimeoutError:
+                level = logging.ERROR if self._seen_non_close_message else logging.DEBUG
+                _LOGGER.log(level, "Websocket timeout: %s", url)
             except Exception:
-                _LOGGER.exception("Error processing websocket message")
+                _LOGGER.debug(
+                    "Unexpected 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:
-        url = self.get_url()
+    async def _websocket_inner_loop(self, url: URL) -> None:
         _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
+        self._seen_non_close_message = False
+        session = await self._get_session()
         # catch any and all errors for Websocket so we can clean up correctly
         try:
             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):
+            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
+                elif msg_type in _CLOSE_MESSAGE_TYPES:
+                    _LOGGER.debug("Websocket closed: %s", msg)
                     break
-                self._reset_timeout()
-        except ClientError:
-            _LOGGER.exception("Websocket disconnect error: %s", url)
+
+                self._seen_non_close_message = True
+                try:
+                    self._subscription(msg)
+                except Exception:
+                    _LOGGER.exception("Error processing websocket message")
         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 self._seen_non_close_message
+                    or (msg.extra and "lastUpdateId" in msg.extra)
+                )
+            ):
+                self._update_bootstrap()
+            _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
+    async def _attempt_reauth(self) -> None:
+        """Attempt to re-authenticate."""
         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")
+            self._headers = await self._auth(True)
+        except Exception:
+            _LOGGER.exception("Error reauthenticating websocket")
 
-        start_event = asyncio.Event()
-        _LOGGER.debug("Scheduling WS connect...")
-        self._websocket_loop_task = asyncio.create_task(
-            self._websocket_loop(start_event),
-        )
+    def start(self) -> None:
+        """Start the websocket."""
+        if self._running:
+            return
+        self._running = True
+        self._websocket_loop_task = asyncio.create_task(self._websocket_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
+            self._stop_task = None
+
+    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.20.0.dist-info → uiprotect-2.1.0.dist-info}/METADATA

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

{uiprotect-1.20.0.dist-info → uiprotect-2.1.0.dist-info}/RECORD

@@ -1,7 +1,7 @@
 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=P_iS5nIxc4sAYp4ISBS_X7wUGIPV2c96xhUkj6kGleg,67379
+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=3SJFF8qs1Jz8t3mD8qwc1hFSocolFjdXI_v4yVlC7o4,20088
-uiprotect/websocket.py,sha256=JHI_2EZeRPqPyQopsBZS0dr3tu0HaTiqeLazfBXhW_8,7339
-uiprotect-1.20.0.dist-info/LICENSE,sha256=INx18jhdbVXMEiiBANeKEbrbz57ckgzxk5uutmmcxGk,1111
-uiprotect-1.20.0.dist-info/METADATA,sha256=2tR-i6s0-c1E1PzpEfVbFq9V9jA2RjilNdhywYM9NW0,10983
-uiprotect-1.20.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-uiprotect-1.20.0.dist-info/entry_points.txt,sha256=J78AUTPrTTxgI3s7SVgrmGqDP7piX2wuuEORzhDdVRA,47
-uiprotect-1.20.0.dist-info/RECORD,,
+uiprotect/websocket.py,sha256=qVx2ZXOacnvyipiPy76i19Rpu596hFxGHhYWuMJVOtg,6898
+uiprotect-2.1.0.dist-info/LICENSE,sha256=INx18jhdbVXMEiiBANeKEbrbz57ckgzxk5uutmmcxGk,1111
+uiprotect-2.1.0.dist-info/METADATA,sha256=2ObhpC8nRihfVKuDZlBYh84HQX93a9FkHPRswCTcn44,10982
+uiprotect-2.1.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+uiprotect-2.1.0.dist-info/entry_points.txt,sha256=J78AUTPrTTxgI3s7SVgrmGqDP7piX2wuuEORzhDdVRA,47
+uiprotect-2.1.0.dist-info/RECORD,,