feat: add linked-hub validation and mixed endpoint warnings (#822)

* feat: add linked-hub validation and mixed endpoint warnings

* precise hostname check

---------

Co-authored-by: Yuelin Zhao <yuelinzhao@microsoft.com>

Author: cheatsheet1999

azext_iot/core/_params.py

@@ -147,12 +147,12 @@ def load_arguments(self, _):  # pylint: disable=too-many-statements
                    'Required when authentication type is UserAssigned.')
         c.argument('hostname_type',
                    options_list=['--hostname-type', '--ht'],
-                   arg_type=get_enum_type(["device", "classic"]),
-                   default="device",
+                   arg_type=get_enum_type(["auto", "device", "classic"]),
+                   default="auto",
                    help="Type of IoT Hub hostname to use when linking. "
-                   "'device' uses the TLS 1.3 device hostname (hub.device.azure-devices.net). "
+                   "'auto' uses the TLS 1.3 device hostname if available, classic otherwise. "
+                   "'device' uses the TLS 1.3 device hostname (errors if not GWv2). "
                    "'classic' uses the classic hostname (hub.azure-devices.net). "
-                   "Falls back to classic if the hub does not support TLS 1.3. "
                    "Only applies when --hub-name is provided.")
         c.argument('apply_allocation_policy',
                    help='A boolean indicating whether to apply allocation policy to the IoT hub.',

azext_iot/core/custom.py

@@ -82,11 +82,48 @@ def _get_resource_group_from_hub(hub):
     return hub["resourcegroup"]
 
 
-def _resolve_linked_hub_hostname(hub, hostname_type="device"):
+def _resolve_linked_hub_hostname(hub, hostname_type="auto"):
     """Resolve IoT Hub hostname for DPS linked hub based on hostname type."""
-    if hostname_type == "device":
-        return hub["properties"].get("deviceHostName") or hub["properties"]["hostName"]
-    return hub["properties"]["hostName"]
+    if hostname_type == "classic":
+        return hub["properties"]["hostName"]
+    device_hostname = hub["properties"].get("deviceHostName")
+    if hostname_type == "device" and not device_hostname:
+        hub_name = hub.get("name", "unknown")
+        raise InvalidArgumentValueError(
+            f"The device hostname is not available for IoT Hub '{hub_name}'. "
+            "This hostname type is only supported on GWv2 IoT Hubs. "
+            "Use '--hostname-type classic' or '--hostname-type auto' instead."
+        )
+    # "auto" or "device" with available deviceHostName
+    return device_hostname or hub["properties"]["hostName"]
+
+
+def _warn_mixed_endpoint_types(linked_hubs):
+    """Warn if DPS dynamic allocation references hubs with mixed hostname types."""
+    types = set()
+    for hub in linked_hubs:
+        # Only check hubs participating in allocation
+        if hub.get("applyAllocationPolicy") is False:
+            continue
+        hostname = hub.get("hostName", "")
+        if not hostname:
+            cs = hub.get("connectionString", "")
+            for part in cs.split(";"):
+                if part.lower().startswith("hostname="):
+                    hostname = part.split("=", 1)[1]
+                    break
+        if not hostname:
+            hostname = hub.get("name", "")
+        parts = hostname.split(".")
+        if len(parts) > 1 and parts[1] == "device":
+            types.add("device")
+        elif hostname:
+            types.add("classic")
+    if len(types) > 1:
+        logger.warning(
+            "DPS has linked hubs with mixed hostname types (device and classic). "
+            "This may cause inconsistent behavior during device provisioning."
+        )
 
 
 # CUSTOM METHODS FOR DPS
@@ -340,7 +377,7 @@ def iot_dps_linked_hub_create(
     resource_group_name=None,
     authentication_type=None,
     user_assigned_identity=None,
-    hostname_type="device",
+    hostname_type="auto",
     apply_allocation_policy=None,
     allocation_weight=None,
     no_wait=False
@@ -443,6 +480,9 @@ def iot_dps_linked_hub_create(
 
     dps["properties"]["iotHubs"].append(linked_hub_entry)
 
+    # Warn if linked hubs have mixed hostname types (device + classic)
+    _warn_mixed_endpoint_types(dps["properties"]["iotHubs"])
+
     if no_wait:
         return client.iot_dps_resource.begin_create_or_update(
             resource_group_name=resource_group_name, provisioning_service_name=dps_name, iot_dps_description=dps

azext_iot/tests/dps/core/test_dps_linked_hub_unit.py

@@ -10,25 +10,36 @@
     MutuallyExclusiveArgumentError,
     RequiredArgumentMissingError,
 )
-from azext_iot.core.custom import _resolve_linked_hub_hostname
+from azext_iot.core.custom import _resolve_linked_hub_hostname, _warn_mixed_endpoint_types
 
 
 class TestResolveLinkedHubHostname:
     def test_device_with_tls13(self):
         hub = {"properties": {"deviceHostName": "hub.device.azure-devices.net", "hostName": "hub.azure-devices.net"}}
         assert _resolve_linked_hub_hostname(hub, "device") == "hub.device.azure-devices.net"
 
-    def test_device_fallback_to_classic(self):
+    def test_device_errors_on_v1_hub(self):
+        hub = {"properties": {"hostName": "hub.azure-devices.net"}, "name": "hub"}
+        with pytest.raises(InvalidArgumentValueError, match="device hostname is not available"):
+            _resolve_linked_hub_hostname(hub, "device")
+
+    def test_auto_fallback_to_classic(self):
         hub = {"properties": {"hostName": "hub.azure-devices.net"}}
-        assert _resolve_linked_hub_hostname(hub, "device") == "hub.azure-devices.net"
+        assert _resolve_linked_hub_hostname(hub, "auto") == "hub.azure-devices.net"
+
+    def test_auto_uses_device_when_available(self):
+        hub = {"properties": {"deviceHostName": "hub.device.azure-devices.net", "hostName": "hub.azure-devices.net"}}
+        assert _resolve_linked_hub_hostname(hub, "auto") == "hub.device.azure-devices.net"
 
     def test_classic(self):
         hub = {"properties": {"deviceHostName": "hub.device.azure-devices.net", "hostName": "hub.azure-devices.net"}}
         assert _resolve_linked_hub_hostname(hub, "classic") == "hub.azure-devices.net"
 
-    def test_default_is_device(self):
+    def test_default_is_auto(self):
         hub = {"properties": {"deviceHostName": "hub.device.azure-devices.net", "hostName": "hub.azure-devices.net"}}
         assert _resolve_linked_hub_hostname(hub) == "hub.device.azure-devices.net"
+        hub_v1 = {"properties": {"hostName": "hub.azure-devices.net"}}
+        assert _resolve_linked_hub_hostname(hub_v1) == "hub.azure-devices.net"
 
 
 class TestLinkedHubCreateValidation:
@@ -111,3 +122,50 @@ def test_mi_null_identity_on_dps(self, fixture_cmd, mock_deps, mocker):
                 cmd=fixture_cmd, client=mock_deps, dps_name="dps",
                 hub_name="hub", authentication_type="SystemAssigned"
             )
+
+
+class TestMixedEndpointWarning:
+    def test_no_warning_all_device(self, caplog):
+        hubs = [
+            {"name": "hub1.device.azure-devices.net"},
+            {"name": "hub2.device.azure-devices.net"},
+        ]
+        _warn_mixed_endpoint_types(hubs)
+        assert "mixed hostname types" not in caplog.text
+
+    def test_no_warning_all_classic(self, caplog):
+        hubs = [
+            {"name": "hub1.azure-devices.net"},
+            {"name": "hub2.azure-devices.net"},
+        ]
+        _warn_mixed_endpoint_types(hubs)
+        assert "mixed hostname types" not in caplog.text
+
+    def test_warning_on_mixed(self, caplog):
+        import logging
+        with caplog.at_level(logging.WARNING):
+            hubs = [
+                {"name": "hub1.device.azure-devices.net"},
+                {"name": "hub2.azure-devices.net"},
+            ]
+            _warn_mixed_endpoint_types(hubs)
+            assert "mixed hostname types" in caplog.text
+
+    def test_warning_on_mixed_with_connection_string(self, caplog):
+        import logging
+        with caplog.at_level(logging.WARNING):
+            hubs = [
+                {"name": "hub1.device.azure-devices.net"},
+                {"connectionString": "HostName=hub2.azure-devices.net;SharedAccessKeyName=x;SharedAccessKey=y"},
+            ]
+            _warn_mixed_endpoint_types(hubs)
+            assert "mixed hostname types" in caplog.text
+
+    def test_no_warning_single_hub(self, caplog):
+        hubs = [{"name": "hub1.device.azure-devices.net"}]
+        _warn_mixed_endpoint_types(hubs)
+        assert "mixed hostname types" not in caplog.text
+
+    def test_no_warning_empty(self, caplog):
+        _warn_mixed_endpoint_types([])
+        assert "mixed hostname types" not in caplog.text