pylxd/client: Support ClientConnection on events

This handles the new default `_WebsocketClient` appropriately and accepts a custom `ClientConnection` as an argument. If provided client is a `ws4py` client instead, fallback on handling it as such for backwards compatibility.

Signed-off-by: hamistao <[email protected]>

pylxd/client.py

@@ -16,6 +16,8 @@
 import os
 import re
 import socket
+import ssl
+import warnings
 from enum import Enum
 from typing import NamedTuple
 from urllib import parse
@@ -670,7 +672,7 @@ def events(self, websocket_client=None, event_types=None):
 
         :param websocket_client: Optional websocket client can be specified for
          implementation-specific handling of events as they occur.
-        :type websocket_client: ws4py.client import WebSocketBaseClient
+        :type websocket_client: websockets.sync.client.ClientConnection or ws4py.client.WebSocketBaseClient
 
         :param event_types: Optional set of event types to propagate. Omit this
          argument or specify {EventTypes.All} to receive all events.
@@ -683,10 +685,6 @@ def events(self, websocket_client=None, event_types=None):
             websocket_client = _WebsocketClient
 
         use_ssl = self.api.scheme == "https" and self.cert
-        ssl_options = (
-            {"certfile": self.cert[0], "keyfile": self.cert[1]} if use_ssl else None
-        )
-        client = websocket_client(self.websocket_url, ssl_options=ssl_options)
         parsed = parse.urlparse(self.api.events._api_endpoint)
 
         resource = parsed.path
@@ -696,6 +694,56 @@ def events(self, websocket_client=None, event_types=None):
             query.update({"type": ",".join(t.value for t in event_types)})
             resource = f"{resource}?{parse.urlencode(query)}"
 
-        client.resource = resource
+        # First try handling it as a websockets' ClientConnection
+        if issubclass(websocket_client, ClientConnection):
+            ssl_context = (
+                ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT).load_cert_chain(
+                    certfile=self.cert[0], keyfile=self.cert[1]
+                )
+                if use_ssl
+                else None
+            )
+
+            is_unix_socket = "+unix" in parsed.scheme
+
+            return create_client_connection(
+                websocket_client=websocket_client,
+                websocket_url=self.websocket_url,
+                unix_socket_path=(
+                    parsed.hostname.replace(
+                        "%2F", "/"
+                    )  # urlparse fails to make sense of slashes in the hostname part.
+                    if is_unix_socket
+                    else None
+                ),
+                resource=resource,
+                ssl=ssl_context,
+            )
+
+        # If not a ClientConnection, assume it is a ws4py client and handle it accordingly, for backwards compatibility.
+        else:
+            try:
+                ssl_options = (
+                    {"certfile": self.cert[0], "keyfile": self.cert[1]}
+                    if use_ssl
+                    else None
+                )
+
+                client = websocket_client(
+                    self.websocket_url, resource, ssl_options=ssl_options
+                )
+
+                client.resource = resource
+
+                warnings.warn(
+                    "The ws4py client API is deprecated and should not be supported in the future",
+                    DeprecationWarning,
+                )
+            except (TypeError, ValueError, AttributeError):
+                warnings.warn(
+                    "Could not create client object since provided client follows neither websockets' nor ws4py's client APIs"
+                )
+
+                return None
 
         return client