From d670e79e49d36d700a4383bdadfad51de1723f2c Mon Sep 17 00:00:00 2001 From: James Garner Date: Wed, 12 Mar 2025 13:01:46 +1300 Subject: [PATCH 1/3] fix: create websockets ConnectionClosed error correctly --- juju/model/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/juju/model/__init__.py b/juju/model/__init__.py index 826cff88..7e926fcb 100644 --- a/juju/model/__init__.py +++ b/juju/model/__init__.py @@ -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.ConnectionClosed() @property def tag(self): From ec3420f7dc03926e0a388c7c0059366bca67a902 Mon Sep 17 00:00:00 2001 From: James Garner Date: Wed, 12 Mar 2025 14:32:30 +1300 Subject: [PATCH 2/3] test: update mock websocket's use of ConnectionClosed --- tests/unit/test_connection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_connection.py b/tests/unit/test_connection.py index 9ed87646..02c90ff4 100644 --- a/tests/unit/test_connection.py +++ b/tests/unit/test_connection.py @@ -26,7 +26,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 ConnectionClosed(None, "ran out of responses") return json.dumps(self.responses.popleft()) async def close(self): From 416de028a2e205cbf0a09cc24b9f2c7dec0d4297 Mon Sep 17 00:00:00 2001 From: James Garner Date: Tue, 18 Mar 2025 11:44:50 +1300 Subject: [PATCH 3/3] fix: raise ConnectionClosed subtypes, correctly (+ update test raise) Two arguments are required, rcvd and sent, both of which should be a `websockets.frames.Closed` object or None. If neither are None, the third argument should also be provided to say in which order sent and rcvd occurred. The existing style appears to have been based on an older API which looks more similar to raising a `Closed` object directly (a code and reason are provided). The exception types are exposed on websockets directly, so I've also removed references to websockets.exceptions.ConnectionClosed. --- juju/client/connection.py | 17 +++++++---------- juju/model/__init__.py | 4 ++-- tests/unit/test_connection.py | 3 +-- 3 files changed, 10 insertions(+), 14 deletions(-) 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 7e926fcb..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() + 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 02c90ff4..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(None, "ran out of responses") + raise websockets.ConnectionClosedOK(rvcd=None, sent=None) return json.dumps(self.responses.popleft()) async def close(self):