fix: ensure IoT Hub list persists over DPS enrollment (group) updates

azext_iot/_help.py

@@ -1114,7 +1114,7 @@
       text: >
         az iot dps enrollment create -g {resource_group_name} --dps-name {dps_name}
         --enrollment-id {enrollment_id} --attestation-type tpm --allocation-policy hashed
-        --endorsement-key 14963E8F3BA5B3984110B3C1CA8E8B89 --iot-hubs "{iot_hub_host_name1} {iot_hub_host_name2}"
+        --endorsement-key 14963E8F3BA5B3984110B3C1CA8E8B89 --iot-hubs {iot_hub_host_name1} {iot_hub_host_name2}
     - name: Create an enrollment 'MyEnrollment' with custom allocation policy,
       text: >
         az iot dps enrollment create -g {resource_group_name} --dps-name {dps_name}
@@ -1161,7 +1161,7 @@
       text: >
         az iot dps enrollment update -g {resource_group_name} --dps-name {dps_name}
         --enrollment-id {enrollment_id} --allocation-policy geolatency
-        --etag AAAAAAAAAAA= --iot-hubs "{iot_hub_host_name1} {iot_hub_host_name2} {iot_hub_host_name3}"
+        --etag AAAAAAAAAAA= --iot-hubs {iot_hub_host_name1} {iot_hub_host_name2} {iot_hub_host_name3}
     - name: Update enrollment '{enrollment_id}' in the Azure IoT Device Provisioning Service '{dps_name}'
             in the resource group '{resource_group_name}' with
             initial twin properties '{"location":{"region":"USA"}}', initial twin tags '{"version":"2"}',

azext_iot/_params.py

@@ -1006,6 +1006,8 @@ def load_arguments(self, _):
             options_list=["--iot-hubs", "--ih"],
             help="Host name of target IoT Hub associated with the allocation policy. Use space-separated "
             "list for multiple IoT Hubs.",
+            nargs="+",
+            action="extend",
             arg_group="Allocation Policy"
         )
         context.argument(

azext_iot/operations/dps.py

@@ -181,7 +181,7 @@ def iot_dps_device_enrollment_create(
             )
         reprovision = _get_reprovision_policy(reprovision_policy)
         initial_twin = _get_initial_twin(initial_twin_tags, initial_twin_properties)
-        iot_hub_list = iot_hubs.split() if iot_hubs else iot_hubs
+        iot_hub_list = iot_hubs.split() if isinstance(iot_hubs, str) else iot_hubs
         _validate_allocation_policy_for_enrollment(
             allocation_policy, iot_hub_host_name, iot_hub_list, webhook_url, api_version
         )
@@ -303,18 +303,27 @@ def iot_dps_device_enrollment_update(
         enrollment_record.initial_twin = _get_updated_inital_twin(
             enrollment_record, initial_twin_tags, initial_twin_properties
         )
-        iot_hub_list = iot_hubs.split() if iot_hubs else iot_hubs
+        iot_hub_list = iot_hubs.split() if isinstance(iot_hubs, str) else iot_hubs
         _validate_allocation_policy_for_enrollment(
-            allocation_policy, iot_hub_host_name, iot_hub_list, webhook_url, api_version
+            allocation_policy,
+            iot_hub_host_name,
+            iot_hub_list,
+            webhook_url,
+            api_version,
+            current_enrollment=enrollment_record
         )
-        if allocation_policy:
-            enrollment_record.allocation_policy = allocation_policy
+        if iot_hub_list:
             enrollment_record.iot_hubs = iot_hub_list
             enrollment_record.iot_hub_host_name = None
-            if allocation_policy == AllocationType.custom.value:
-                enrollment_record.custom_allocation_definition = CustomAllocationDefinition(
-                    webhook_url=webhook_url, api_version=api_version
-                )
+        if allocation_policy:
+            enrollment_record.allocation_policy = allocation_policy
+        if enrollment_record.allocation_policy == AllocationType.custom.value and any([
+            webhook_url, api_version
+        ]):
+            enrollment_record.custom_allocation_definition = CustomAllocationDefinition(
+                webhook_url=webhook_url or enrollment_record.custom_allocation_definition.webhook_url,
+                api_version=api_version or enrollment_record.custom_allocation_definition.api_version
+            )
         if edge_enabled is not None:
             enrollment_record.capabilities = DeviceCapabilities(iot_edge=edge_enabled)
         if device_information:
@@ -481,7 +490,7 @@ def iot_dps_device_enrollment_group_create(
             )
         reprovision = _get_reprovision_policy(reprovision_policy)
         initial_twin = _get_initial_twin(initial_twin_tags, initial_twin_properties)
-        iot_hub_list = iot_hubs.split() if iot_hubs else iot_hubs
+        iot_hub_list = iot_hubs.split() if isinstance(iot_hubs, str) else iot_hubs
         _validate_allocation_policy_for_enrollment(
             allocation_policy, iot_hub_host_name, iot_hub_list, webhook_url, api_version
         )
@@ -621,18 +630,27 @@ def iot_dps_device_enrollment_group_update(
         enrollment_record.initial_twin = _get_updated_inital_twin(
             enrollment_record, initial_twin_tags, initial_twin_properties
         )
-        iot_hub_list = iot_hubs.split() if iot_hubs else iot_hubs
+        iot_hub_list = iot_hubs.split() if isinstance(iot_hubs, str) else iot_hubs
         _validate_allocation_policy_for_enrollment(
-            allocation_policy, iot_hub_host_name, iot_hub_list, webhook_url, api_version
+            allocation_policy,
+            iot_hub_host_name,
+            iot_hub_list,
+            webhook_url,
+            api_version,
+            current_enrollment=enrollment_record
         )
-        if allocation_policy:
-            enrollment_record.allocation_policy = allocation_policy
+        if iot_hub_list:
             enrollment_record.iot_hubs = iot_hub_list
             enrollment_record.iot_hub_host_name = None
-            if allocation_policy == AllocationType.custom.value:
-                enrollment_record.custom_allocation_definition = CustomAllocationDefinition(
-                    webhook_url=webhook_url, api_version=api_version
-                )
+        if allocation_policy:
+            enrollment_record.allocation_policy = allocation_policy
+        if enrollment_record.allocation_policy == AllocationType.custom.value and any([
+            webhook_url, api_version
+        ]):
+            enrollment_record.custom_allocation_definition = CustomAllocationDefinition(
+                webhook_url=webhook_url or enrollment_record.custom_allocation_definition.webhook_url,
+                api_version=api_version or enrollment_record.custom_allocation_definition.api_version
+            )
         if edge_enabled is not None:
             enrollment_record.capabilities = DeviceCapabilities(iot_edge=edge_enabled)
         return sdk.enrollment_group.create_or_update(
@@ -1137,8 +1155,16 @@ def _validate_arguments_for_attestation_mechanism(
 
 
 def _validate_allocation_policy_for_enrollment(
-    allocation_policy, iot_hub_host_name, iot_hub_list, webhook_url, api_version
+    allocation_policy, iot_hub_host_name, iot_hub_list, webhook_url, api_version, current_enrollment=None
 ):
+    # get the enrollment values if not provided but present
+    if current_enrollment:
+        iot_hub_list = iot_hub_list or current_enrollment.iot_hubs
+        allocation_policy = allocation_policy or current_enrollment.allocation_policy
+        if current_enrollment.allocation_policy == AllocationType.custom.value:
+            webhook_url = webhook_url or current_enrollment.custom_allocation_definition.webhook_url
+            api_version = api_version or current_enrollment.custom_allocation_definition.api_version
+
     if allocation_policy:
         if iot_hub_host_name is not None:
             raise MutuallyExclusiveArgumentError(
@@ -1149,6 +1175,7 @@ def _validate_allocation_policy_for_enrollment(
             allocation_policy == allocation.value for allocation in AllocationType
         ):
             raise RequiredArgumentMissingError("Please provide valid allocation policy.")
+
         if allocation_policy == AllocationType.static.value:
             if iot_hub_list is None:
                 raise RequiredArgumentMissingError("Please provide a hub to be assigned with device.")
@@ -1160,6 +1187,5 @@ def _validate_allocation_policy_for_enrollment(
                     "Please provide both the Azure function webhook url and provisioning"
                     " service api-version when the allocation-policy is defined as Custom."
                 )
-    else:
-        if iot_hub_list:
-            raise RequiredArgumentMissingError("Please provide allocation policy.")
+    elif iot_hub_list and not current_enrollment:
+        raise RequiredArgumentMissingError("Please provide allocation policy.")

azext_iot/tests/dps/enrollment/test_iot_dps_enrollment_int.py

@@ -298,7 +298,7 @@ def test_dps_enrollment_symmetrickey_lifecycle(provisioned_iot_dps_module):
         assert update_enrollment["customAllocationDefinition"]["webhookUrl"] == WEBHOOK_URL
         assert update_enrollment["customAllocationDefinition"]["apiVersion"] == API_VERSION
         assert update_enrollment["deviceId"] == device_id
-        assert update_enrollment["iotHubs"] is None
+        assert update_enrollment["iotHubs"] == [hub_hostname]
         assert update_enrollment["initialTwin"]["tags"]
         assert update_enrollment["initialTwin"]["properties"]["desired"]
         assert update_enrollment["optionalDeviceInformation"] == generic_dict

azext_iot/tests/dps/enrollment/test_iot_dps_enrollment_unit.py

@@ -123,6 +123,12 @@ def serviceclient_generic_error(self, mocked_response, fixture_gdcs, fixture_dps
                                         reprovision_policy='never',
                                         allocation_policy='hashed',
                                         iot_hubs='hub1 hub2')),
+        (generate_enrollment_create_req(attestation_type='symmetricKey',
+                                        primary_key='primarykey',
+                                        secondary_key='secondarykey',
+                                        reprovision_policy='never',
+                                        allocation_policy='hashed',
+                                        iot_hubs=['hub1', 'hub2'])),
         (generate_enrollment_create_req(attestation_type='symmetricKey',
                                         primary_key='primarykey',
                                         secondary_key='secondarykey',
@@ -227,7 +233,9 @@ def test_enrollment_create(self, serviceclient, fixture_cmd, req):
                 assert body['customAllocationDefinition']['webhookUrl'] == req['webhook_url']
                 assert body['customAllocationDefinition']['apiVersion'] == req['api_version']
         if req['iot_hubs']:
-            assert body['iotHubs'] == req['iot_hubs'].split()
+            assert body['iotHubs'] == (
+                req['iot_hubs'].split() if isinstance(req['iot_hubs'], str) else req['iot_hubs']
+            )
         if req['edge_enabled']:
             assert body['capabilities']['iotEdge']
 
@@ -376,14 +384,16 @@ def serviceclient(self, mocked_response, fixture_gdcs, fixture_dps_sas, patch_ce
         (generate_enrollment_update_req(reprovision_policy='never')),
         (generate_enrollment_update_req(allocation_policy='static', iot_hubs='hub1')),
         (generate_enrollment_update_req(allocation_policy='hashed', iot_hubs='hub1 hub2')),
+        (generate_enrollment_update_req(allocation_policy='hashed', iot_hubs=['hub1', 'hub2'])),
         (generate_enrollment_update_req(allocation_policy='geoLatency')),
         (generate_enrollment_update_req(allocation_policy='custom',
                                         webhook_url="https://www.test.test",
                                         api_version="2019-03-31")),
         (generate_enrollment_update_req(edge_enabled=True)),
         (generate_enrollment_update_req(edge_enabled=False))
     ])
-    def test_enrollment_update(self, serviceclient, fixture_cmd, req):
+    def test_enrollment_update(self, mocker, serviceclient, fixture_cmd, req):
+        mocker.patch("azext_iot.operations.dps._validate_allocation_policy_for_enrollment")
         subject.iot_dps_device_enrollment_update(
             cmd=fixture_cmd,
             enrollment_id=req['enrollment_id'],
@@ -457,7 +467,9 @@ def test_enrollment_update(self, serviceclient, fixture_cmd, req):
                 assert body['customAllocationDefinition']['webhookUrl'] == req['webhook_url']
                 assert body['customAllocationDefinition']['apiVersion'] == req['api_version']
         if req['iot_hubs']:
-            assert body['iotHubs'] == req['iot_hubs'].split()
+            assert body['iotHubs'] == (
+                req['iot_hubs'].split() if isinstance(req['iot_hubs'], str) else req['iot_hubs']
+            )
         if req['edge_enabled'] is not None:
             assert body['capabilities']['iotEdge'] == req['edge_enabled']
 

azext_iot/tests/dps/enrollment_group/test_iot_dps_enrollment_group_int.py

@@ -248,7 +248,7 @@ def test_dps_enrollment_group_symmetrickey_lifecycle(provisioned_iot_dps_module)
         assert enrollment_update["customAllocationDefinition"]["webhookUrl"] == WEBHOOK_URL
         assert enrollment_update["customAllocationDefinition"]["apiVersion"] == API_VERSION
         assert enrollment_update["enrollmentGroupId"] == enrollment_id
-        assert enrollment_update["iotHubs"] is None
+        assert enrollment_update["iotHubs"] == [hub_hostname]
         assert enrollment_update["initialTwin"]["tags"]
         assert enrollment_update["initialTwin"]["properties"]["desired"]
         assert enrollment_update["provisioningStatus"] == EntityStatusType.disabled.value

azext_iot/tests/dps/enrollment_group/test_iot_dps_enrollment_group_unit.py

@@ -96,6 +96,9 @@ def serviceclient(self, mocked_response, fixture_gdcs, fixture_dps_sas, patch_ce
         (generate_enrollment_group_create_req(certificate_path='myCert',
                                               allocation_policy='hashed',
                                               iot_hubs='hub1 hub2')),
+        (generate_enrollment_group_create_req(certificate_path='myCert',
+                                              allocation_policy='hashed',
+                                              iot_hubs=['hub1', 'hub2'])),
         (generate_enrollment_group_create_req(certificate_path='myCert',
                                               allocation_policy='geoLatency')),
         (generate_enrollment_group_create_req(certificate_path='myCert',
@@ -190,7 +193,9 @@ def test_enrollment_group_create(self, serviceclient, fixture_cmd, req):
                 assert body['customAllocationDefinition']['webhookUrl'] == req['webhook_url']
                 assert body['customAllocationDefinition']['apiVersion'] == req['api_version']
         if req['iot_hubs']:
-            assert body['iotHubs'] == req['iot_hubs'].split()
+            assert body['iotHubs'] == ((
+                req['iot_hubs'].split() if isinstance(req['iot_hubs'], str) else req['iot_hubs'])
+            )
         if req['edge_enabled']:
             assert body['capabilities']['iotEdge']
 
@@ -362,12 +367,14 @@ def serviceclient(self, mocked_response, fixture_gdcs, fixture_dps_sas, patch_ce
                                               webhook_url="https://www.test.test",
                                               api_version="2019-03-31")),
         (generate_enrollment_group_update_req(allocation_policy='hashed', iot_hubs='hub1 hub2')),
+        (generate_enrollment_group_update_req(allocation_policy='hashed', iot_hubs=['hub1', 'hub2'])),
         (generate_enrollment_group_update_req(allocation_policy='geoLatency')),
         (generate_enrollment_group_update_req(iot_hub_host_name='hub1')),
         (generate_enrollment_group_update_req(edge_enabled=True)),
         (generate_enrollment_group_update_req(edge_enabled=False))
     ])
-    def test_enrollment_group_update(self, serviceclient, fixture_cmd, req):
+    def test_enrollment_group_update(self, mocker, serviceclient, fixture_cmd, req):
+        mocker.patch("azext_iot.operations.dps._validate_allocation_policy_for_enrollment")
         subject.iot_dps_device_enrollment_group_update(
             cmd=fixture_cmd,
             enrollment_id=req['enrollment_id'],
@@ -445,7 +452,9 @@ def test_enrollment_group_update(self, serviceclient, fixture_cmd, req):
                 assert body['customAllocationDefinition']['webhookUrl'] == req['webhook_url']
                 assert body['customAllocationDefinition']['apiVersion'] == req['api_version']
         if req['iot_hubs']:
-            assert body['iotHubs'] == req['iot_hubs'].split()
+            assert body['iotHubs'] == ((
+                req['iot_hubs'].split() if isinstance(req['iot_hubs'], str) else req['iot_hubs'])
+            )
         if req['edge_enabled'] is not None:
             assert body['capabilities']['iotEdge'] == req['edge_enabled']