diff --git a/juju/client/connection.py b/juju/client/connection.py index 0f5c2d3e..8337c4fd 100644 --- a/juju/client/connection.py +++ b/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 diff --git a/juju/model/__init__.py b/juju/model/__init__.py index 826cff88..56c3aec4 100644 --- a/juju/model/__init__.py +++ b/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): diff --git a/tests/unit/test_connection.py b/tests/unit/test_connection.py index 9ed87646..8eb6f269 100644 --- a/tests/unit/test_connection.py +++ b/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):