fix: create websockets ConnectionClosed error correctly

juju/client/connection.py

@@ -368,7 +368,7 @@ async def close(self, to_reconnect: bool = False):
                 await asyncio.gather(*tasks_need_to_be_gathered)
             except asyncio.CancelledError:
                 pass
-            except websockets.exceptions.ConnectionClosed:
+            except websockets.ConnectionClosed:
                 pass
 
         self._pinger_task = None
@@ -380,7 +380,7 @@ async def close(self, to_reconnect: bool = False):
 
     async def _recv(self, request_id: int) -> dict[str, Any]:
         if not self.is_open:
-            raise websockets.exceptions.ConnectionClosedOK(None, None)
+            raise websockets.ConnectionClosedOK(rcvd=None, sent=None)
         try:
             return await self.messages.get(request_id)
         except GeneratorExit:
@@ -463,7 +463,7 @@ async def _debug_logger(self):
         except asyncio.CancelledError:
             asyncio.create_task(self.close(), name="Task_Close")  # noqa: RUF006
             raise
-        except websockets.exceptions.ConnectionClosed:
+        except websockets.ConnectionClosed:
             log.warning("Debug Logger: Connection closed, reconnecting")
             # the reconnect has to be done as a task because the receiver will
             # be cancelled by the reconnect and we don't want the reconnect
@@ -489,7 +489,7 @@ async def _receiver(self):
         except asyncio.CancelledError:
             log.debug("Receiver: Cancelled")
             pass
-        except websockets.exceptions.ConnectionClosed as e:
+        except websockets.ConnectionClosed as e:
             log.warning("Receiver: Connection closed, reconnecting")
             await self.messages.put_all(e)
             # the reconnect has to be done as a task because the receiver will
@@ -530,7 +530,7 @@ async def _do_ping():
         except asyncio.CancelledError:
             log.debug("Pinger: Cancelled")
             pass
-        except websockets.exceptions.ConnectionClosed:
+        except websockets.ConnectionClosed:
             # The connection has closed - we can't do anything
             # more until the connection is restarted.
             log.debug("ping failed because of closed connection")
@@ -568,11 +568,8 @@ async def rpc(
         for attempt in range(3):
             if self.monitor.status == Monitor.DISCONNECTED:
                 # closed cleanly; shouldn't try to reconnect
-                raise websockets.exceptions.ConnectionClosed(
-                    websockets.frames.Close(
-                        websockets.frames.CloseCode.NORMAL_CLOSURE, "websocket closed"
-                    )
-                )
+                # we're trying to use a reference to a monitor that was cleanly closed earlier
+                raise websockets.ConnectionClosedOK(rcvd=None, sent=None)
             try:
                 await self._ws.send(outgoing)
                 break

juju/model/__init__.py

@@ -1158,7 +1158,7 @@ async def remove_application(
     async def block_until(self, *conditions, timeout=None, wait_period=0.5):
         """Return only after all conditions are true.
 
-        Raises `websockets.ConnectionClosed` if disconnected.
+        Raises `websockets.ConnectionClosedError` if disconnected.
         """
 
         def _disconnected():
@@ -1169,7 +1169,7 @@ def done():
 
         await utils.block_until(done, timeout=timeout, wait_period=wait_period)
         if _disconnected():
-            raise websockets.ConnectionClosed(1006, "no reason")
+            raise websockets.ConnectionClosedError(rcvd=None, sent=None)
 
     @property
     def tag(self):

tests/unit/test_connection.py

@@ -8,7 +8,6 @@
 
 import pytest
 import websockets
-from websockets.exceptions import ConnectionClosed
 
 from juju.client.connection import Connection
 from juju.errors import JujuRedirectException
@@ -26,7 +25,7 @@ async def send(self, message):
     async def recv(self):
         if not self.responses:
             await asyncio.sleep(1)  # delay to give test time to finish
-            raise ConnectionClosed(0, "ran out of responses")
+            raise websockets.ConnectionClosedOK(rvcd=None, sent=None)
         return json.dumps(self.responses.popleft())
 
     async def close(self):