More closing fixes

lexiforest · Dec 1, 2024 · 22fb9db · 22fb9db
1 parent 5b18117
commit 22fb9db
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/curl_cffi/requests/websockets.py b/curl_cffi/requests/websockets.py
@@ -381,7 +381,7 @@ def send(self, payload: Union[str, bytes], flags: CurlWsFlag = CurlWsFlag.BINARY
             flags: flags for the frame.
         """
         if self.closed:
-            raise TypeError("WebSocket is closed")
+            raise SessionClosed("WebSocket is closed")
 
         # curl expects bytes
         if isinstance(payload, str):
@@ -525,7 +525,7 @@ def loop(self):
 
     async def __aiter__(self) -> Self:
         if self.closed:
-            raise TypeError("WebSocket is closed")
+            raise SessionClosed("WebSocket is closed")
         return self
 
     async def __anext__(self) -> bytes:
@@ -541,7 +541,7 @@ async def recv_fragment(self, *, timeout: Optional[float] = None) -> Tuple[bytes
             timeout: how many seconds to wait before giving up.
         """
         if self.closed:
-            raise TypeError("WebSocket is closed")
+            raise SessionClosed("WebSocket is closed")
         if self._recv_lock.locked():
             raise TypeError("Concurrent call to recv_fragment() is not allowed")
 
@@ -551,15 +551,15 @@ async def recv_fragment(self, *, timeout: Optional[float] = None) -> Tuple[bytes
             )
             if frame.flags & CurlWsFlag.CLOSE:
                 try:
-                    self._close_code, self._close_reason = self._unpack_close_frame(chunk)
+                    code, message = self._close_code, self._close_reason = self._unpack_close_frame(chunk)
                 except WebSocketError as e:
                     # Follow the spec to close the connection
                     # Errors do not respect autoclose
                     self._close_code = e.code
                     await self.close(e.code)
                     raise
                 if self.autoclose:
-                    await self.close()
+                    await self.close(code, message.encode())
 
         return chunk, frame
 
@@ -628,7 +628,7 @@ async def send(self, payload: Union[str, bytes], flags: CurlWsFlag = CurlWsFlag.
             flags: flags for the frame.
         """
         if self.closed:
-            raise TypeError("WebSocket is closed")
+            raise SessionClosed("WebSocket is closed")
 
         # curl expects bytes
         if isinstance(payload, str):