Catch and wrap aiohttp and mashumaro errors (#8)

Author: emontnemery

go2rtc_client/exceptions.py

@@ -2,6 +2,45 @@
 
 from __future__ import annotations
 
+from functools import wraps
+from typing import TYPE_CHECKING, Any
+
+from aiohttp import ClientError
+from mashumaro.exceptions import (
+    ExtraKeysError,
+    InvalidFieldValue,
+    MissingDiscriminatorError,
+    MissingField,
+    SuitableVariantNotFoundError,
+    UnserializableDataError,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Coroutine
+
 
 class Go2RtcClientError(Exception):
     """Base exception for go2rtc client."""
+
+
+def handle_error[**_P, _R](
+    func: Callable[_P, Coroutine[Any, Any, _R]],
+) -> Callable[_P, Coroutine[Any, Any, _R]]:
+    """Wrap aiohttp and mashumaro errors."""
+
+    @wraps(func)
+    async def _func(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+        try:
+            return await func(*args, **kwargs)
+        except (
+            ClientError,
+            ExtraKeysError,
+            InvalidFieldValue,
+            MissingDiscriminatorError,
+            MissingField,
+            SuitableVariantNotFoundError,
+            UnserializableDataError,
+        ) as exc:
+            raise Go2RtcClientError from exc
+
+    return _func

go2rtc_client/rest.py

@@ -11,6 +11,7 @@
 from mashumaro.mixins.dict import DataClassDictMixin
 from yarl import URL
 
+from .exceptions import handle_error
 from .models import Stream, WebRTCSdpAnswer, WebRTCSdpOffer
 
 if TYPE_CHECKING:
@@ -78,6 +79,7 @@ async def _forward_sdp_offer(
         )
         return WebRTCSdpAnswer.from_dict(await resp.json())
 
+    @handle_error
     async def forward_whep_sdp_offer(
         self, source_name: str, offer: WebRTCSdpOffer
     ) -> WebRTCSdpAnswer:
@@ -99,11 +101,13 @@ def __init__(self, client: _BaseClient) -> None:
         """Initialize Client."""
         self._client = client
 
+    @handle_error
     async def list(self) -> dict[str, Stream]:
         """List streams registered with the server."""
         resp = await self._client.request("GET", self.PATH)
         return _GET_STREAMS_DECODER.decode(await resp.json())
 
+    @handle_error
     async def add(self, name: str, source: str) -> None:
         """Add a stream to the server."""
         await self._client.request(

go2rtc_client/ws/client.py

@@ -6,16 +6,11 @@
 from typing import TYPE_CHECKING, Any
 from urllib.parse import urljoin
 
-from aiohttp import (
-    ClientError,
-    ClientSession,
-    ClientWebSocketResponse,
-    WSMsgType,
-    WSServerHandshakeError,
-)
+from aiohttp import ClientSession, ClientWebSocketResponse, WSMsgType
 
-from go2rtc_client.exceptions import Go2RtcClientError
-from go2rtc_client.ws.messages import BaseMessage
+from go2rtc_client.exceptions import handle_error
+
+from .messages import BaseMessage
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -56,26 +51,22 @@ def connected(self) -> bool:
         """Return if we're currently connected."""
         return self._client is not None and not self._client.closed
 
+    @handle_error
     async def connect(self) -> None:
         """Connect to device."""
         async with self._connect_lock:
             if self.connected:
                 return
 
             _LOGGER.debug("Trying to connect to %s", self._server_url)
-            try:
-                self._client = await self._session.ws_connect(
-                    urljoin(self._server_url, "/api/ws"), params=self._params
-                )
-            except (
-                WSServerHandshakeError,
-                ClientError,
-            ) as err:
-                raise Go2RtcClientError(err) from err
+            self._client = await self._session.ws_connect(
+                urljoin(self._server_url, "/api/ws"), params=self._params
+            )
 
             self._rx_task = asyncio.create_task(self._receive_messages())
             _LOGGER.info("Connected to %s", self._server_url)
 
+    @handle_error
     async def close(self) -> None:
         """Close connection."""
         if self.connected:
@@ -85,6 +76,7 @@ async def close(self) -> None:
             self._client = None
             await client.close()
 
+    @handle_error
     async def send(self, message: BaseMessage) -> None:
         """Send a message."""
         if not self.connected: