Client: Allow websocket URIs when using set_custom_host()

2 files changed, 8 insertions, 4 deletions

diff --git a/nbxmpp/addresses.py b/nbxmpp/addresses.py
index 7131778..59be915 100644
--- a/nbxmpp/addresses.py
+++ b/nbxmpp/addresses.py
@@ -142,14 +142,18 @@ class ServerAddresses(Observable):
         if address is None:
             return
 
-        host, protocol, type_ = address
+        host_or_uri, protocol, type_ = address
+        if protocol == ConnectionProtocol.WEBSOCKET:
+            host, uri = None, host_or_uri
+        else:
+            host, uri = host_or_uri, None
 
         self._fallback_addresses = []
         self._addresses = [
             ServerAddress(domain=self._domain,
                           service=None,
                           host=host,
-                          uri=None,
+                          uri=uri,
                           protocol=protocol,
                           type=type_,
                           proxy=None)]
diff --git a/nbxmpp/client.py b/nbxmpp/client.py
index 44c192c..bcb159d 100644
--- a/nbxmpp/client.py
+++ b/nbxmpp/client.py
@@ -171,10 +171,10 @@ class Client(Observable):
     def custom_host(self):
         return self._custom_host
 
-    def set_custom_host(self, host, protocol, type_):
+    def set_custom_host(self, host_or_uri, protocol, type_):
         if self._domain is None:
             raise ValueError('Call set_domain() first before set_custom_host()')
-        self._custom_host = (host, protocol, type_)
+        self._custom_host = (host_or_uri, protocol, type_)
 
     def set_accepted_certificates(self, certificates):
         self._accepted_certificates = certificates