start

sdk/servicebus/azure-servicebus/azure/servicebus/_common/_configuration.py

@@ -59,7 +59,7 @@ def __init__(self, **kwargs):
             if self.custom_endpoint_address.find("//") == -1:
                 self.custom_endpoint_address = "sb://" + self.custom_endpoint_address
             endpoint = urlparse(self.custom_endpoint_address)
-            self.transport_type = TransportType.AmqpOverWebsocket
+            self.transport_type = kwargs.get("transport_type") or TransportType.AmqpOverWebsocket
             self.custom_endpoint_hostname = endpoint.hostname
             if amqp_transport.KIND == "pyamqp":
                 self.custom_endpoint_address += "/$servicebus/websocket"

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_connection.py

@@ -136,8 +136,12 @@ def __init__(  # pylint:disable=too-many-locals,too-many-statements
         custom_endpoint = None
         if custom_endpoint_address:
             custom_parsed_url = urlparse(custom_endpoint_address)
-            custom_port = custom_parsed_url.port or WEBSOCKET_PORT
-            custom_endpoint = f"{custom_parsed_url.hostname}:{custom_port}{custom_parsed_url.path}"
+            if transport_type.value == TransportType.Amqp.value:
+                custom_port = custom_parsed_url.port or SECURE_PORT
+                custom_endpoint = f"{custom_parsed_url.hostname}:{custom_port}"
+            else:
+                custom_port = custom_parsed_url.port or WEBSOCKET_PORT
+                custom_endpoint = f"{custom_parsed_url.hostname}:{custom_port}{custom_parsed_url.path}"
         self._container_id = container_id or str(uuid.uuid4())
         self._network_trace = network_trace
         self._network_trace_params = {"amqpConnection": self._container_id, "amqpSession": "", "amqpLink": ""}
@@ -163,6 +167,7 @@ def __init__(  # pylint:disable=too-many-locals,too-many-statements
                 credential=kwargs["sasl_credential"],
                 port=self._port,
                 custom_endpoint=custom_endpoint,
+                custom_port=custom_port,
                 socket_timeout=self._socket_timeout,
                 network_trace_params=self._network_trace_params,
                 **kwargs,

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_transport.py

@@ -179,6 +179,8 @@ def __init__(
         self.raise_on_initial_eintr = raise_on_initial_eintr
         self._read_buffer = BytesIO()
         self.host, self.port = to_host_port(host, port)
+        self._custom_endpoint = kwargs.get("custom_endpoint")
+        self._custom_port = kwargs.get("custom_port")
         self.network_trace_params = kwargs.get("network_trace_params")
 
         self.connect_timeout = connect_timeout
@@ -491,7 +493,9 @@ class SSLTransport(_AbstractTransport):
 
     def __init__(self, host, *, port=AMQPS_PORT, socket_timeout=None, ssl_opts=None, **kwargs):
         self.sslopts = ssl_opts if isinstance(ssl_opts, dict) else {}
-        self.sslopts["server_hostname"] = host
+        self._custom_endpoint = kwargs.get("custom_endpoint")
+        self._custom_port = kwargs.get("custom_port")
+        self.sslopts["server_hostname"] = self._custom_endpoint or host
         self._read_buffer = BytesIO()
         super(SSLTransport, self).__init__(host, port=port, socket_timeout=socket_timeout, **kwargs)