fix: close the socket when connecting an exec websocket fails

ops/pebble.py

@@ -3367,10 +3367,14 @@ def _connect_websocket(self, task_id: str, websocket_id: str) -> _WebSocket:
         # case the Pebble side times out (5s), so this side doesn't hang. See:
         # https://github.com/canonical/operator/issues/1246
         sock.settimeout(self.timeout)
-        sock.connect(self.socket_path)
-        url = self._websocket_url(task_id, websocket_id)
-        ws: _WebSocket = websocket.WebSocket(skip_utf8_validation=True)  # type: ignore
-        ws.connect(url, socket=sock)
+        try:
+            sock.connect(self.socket_path)
+            url = self._websocket_url(task_id, websocket_id)
+            ws: _WebSocket = websocket.WebSocket(skip_utf8_validation=True)  # type: ignore
+            ws.connect(url, socket=sock)
+        except BaseException:
+            sock.close()
+            raise
         # Reset to no timeout so connection can be "long polling" when data is
         # being received.
         sock.settimeout(None)

test/test_pebble.py

@@ -4145,6 +4145,22 @@ def _connect_websocket(self, change_id: str, websocket_id: str):
             client.exec(['foo'])
         assert str(excinfo.value) in 'unexpected error connecting to websockets: conn!'
 
+    def test_connect_websocket_closes_socket_on_error(self, monkeypatch: pytest.MonkeyPatch):
+        sockets: list[socket.socket] = []
+        real_socket = socket.socket
+
+        def fake_socket(*args: typing.Any, **kwargs: typing.Any):
+            sock = real_socket(*args, **kwargs)
+            sockets.append(sock)
+            return sock
+
+        monkeypatch.setattr(pebble.socket, 'socket', fake_socket)
+        client = pebble.Client(socket_path='does_not_exist')
+        with pytest.raises(OSError):
+            client._connect_websocket('1', 'stdio')
+        assert len(sockets) == 1
+        assert sockets[0].fileno() == -1  # Closed.
+
     def test_websocket_send_raises(self, client: MockClient):
         stdio, stderr, _ = self.add_responses(client, '123', 0)
         raised = False