Use library timeout exception class

lexiforest · Dec 2, 2024 · 95f1cfb · 95f1cfb
1 parent 1fa7007
commit 95f1cfb
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/curl_cffi/requests/websockets.py b/curl_cffi/requests/websockets.py
@@ -20,7 +20,7 @@
 )
 
 from .options import set_curl_options
-from .exceptions import SessionClosed
+from .exceptions import SessionClosed, Timeout
 from ..aio import CURL_SOCKET_BAD
 from ..const import CurlECode, CurlOpt, CurlWsFlag, CurlInfo
 from ..curl import Curl, CurlError
@@ -546,9 +546,12 @@ async def recv_fragment(self, *, timeout: Optional[float] = None) -> Tuple[bytes
             raise TypeError("Concurrent call to recv_fragment() is not allowed")
 
         async with self._recv_lock:
-            chunk, frame = await asyncio.wait_for(
-                self.loop.run_in_executor(None, self.curl.ws_recv), timeout
-            )
+            try:
+                chunk, frame = await asyncio.wait_for(
+                    self.loop.run_in_executor(None, self.curl.ws_recv), timeout
+                )
+            except asyncio.TimeoutError:
+                raise Timeout("WebSocket recv_fragment() timed out")
             if frame.flags & CurlWsFlag.CLOSE:
                 try:
                     code, message = self._close_code, self._close_reason = self._unpack_close_frame(chunk)