diff --git a/azext_iot/_help.py b/azext_iot/_help.py index d768a5289..28aaaa239 100644 --- a/azext_iot/_help.py +++ b/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"}', diff --git a/azext_iot/_params.py b/azext_iot/_params.py index 16f50509d..54bf532db 100644 --- a/azext_iot/_params.py +++ b/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( diff --git a/azext_iot/operations/dps.py b/azext_iot/operations/dps.py index 985c70f0a..845d6ce41 100644 --- a/azext_iot/operations/dps.py +++ b/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.") diff --git a/azext_iot/tests/dps/enrollment/test_iot_dps_enrollment_int.py b/azext_iot/tests/dps/enrollment/test_iot_dps_enrollment_int.py index 77eb699a9..6db46d72b 100644 --- a/azext_iot/tests/dps/enrollment/test_iot_dps_enrollment_int.py +++ b/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 diff --git a/azext_iot/tests/dps/enrollment/test_iot_dps_enrollment_unit.py b/azext_iot/tests/dps/enrollment/test_iot_dps_enrollment_unit.py index 291f36aba..d111ca12d 100644 --- a/azext_iot/tests/dps/enrollment/test_iot_dps_enrollment_unit.py +++ b/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,6 +384,7 @@ 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", @@ -383,7 +392,8 @@ def serviceclient(self, mocked_response, fixture_gdcs, fixture_dps_sas, patch_ce (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'] diff --git a/azext_iot/tests/dps/enrollment_group/test_iot_dps_enrollment_group_int.py b/azext_iot/tests/dps/enrollment_group/test_iot_dps_enrollment_group_int.py index 6f92e8029..949889ab7 100644 --- a/azext_iot/tests/dps/enrollment_group/test_iot_dps_enrollment_group_int.py +++ b/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 diff --git a/azext_iot/tests/dps/enrollment_group/test_iot_dps_enrollment_group_unit.py b/azext_iot/tests/dps/enrollment_group/test_iot_dps_enrollment_group_unit.py index 4fc3ad4d9..0377ca370 100644 --- a/azext_iot/tests/dps/enrollment_group/test_iot_dps_enrollment_group_unit.py +++ b/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']