feat: add bulk key regeneration to az iot hub device-identity renew-key and az iot hub module-identity renew-key

.azure-devops/merge.yml

@@ -82,7 +82,7 @@ jobs:
 - job: 'run_unit_tests_macOs'
   dependsOn: ['build_and_publish_azure_iot_cli_ext', 'build_and_publish_azure_cli_test_sdk']
   pool:
-    vmImage: 'macOS-11'
+    vmImage: 'macOS-13'
 
   steps:
   - template: templates/run-tests-parallel.yml

.azure-devops/templates/run-tests-parallel.yml

@@ -63,7 +63,6 @@ steps:
 
   - ${{ if eq(parameters.runUnitTests, 'true') }}:
     - script: |
-        pip freeze
         pytest -vv ${{ parameters.path }} -k "_unit.py" --cov=azext_iot --cov-config .coveragerc --junitxml=junit/test-iotext-unit-${{ parameters.name }}.xml
       displayName: '${{ parameters.name }} unit tests'
       env:

.github/workflows/tox.yml

@@ -20,14 +20,14 @@ jobs:
   unit-test:
     name: Unit test ${{ matrix.py }} - ${{ matrix.os }}
     continue-on-error: ${{ inputs.continue-on-error }}
-    runs-on: ${{ matrix.os }}-latest
+    runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         os:
-          - ubuntu
-          - windows
-          - macos
+          - ubuntu-22.04
+          - windows-2022
+          - macos-13
         py:
           - "3.11"
           - "3.10"

azext_iot/_help.py

@@ -231,7 +231,9 @@
 ] = """
     type: command
     short-summary: Renew target keys of an IoT Hub device with sas authentication.
-    long-summary: Currently etags and key type `swap` are not supported for bulk key regeneration.
+    long-summary: |
+                  Currently etags and key type `swap` are not supported for bulk key regeneration.
+                  Bulk Key regeneration will yeild a different output format from single device key regeneration.
     examples:
       - name: Renew the primary key.
         text: az iot hub device-identity renew-key -d {device_id} -n {iothub_name} --kt primary
@@ -563,7 +565,9 @@
 ] = """
     type: command
     short-summary: Renew target keys of an IoT Hub device module with sas authentication.
-    long-summary: Currently etags and key type `swap` are not supported for bulk key regeneration.
+    long-summary: |
+                  Currently etags and key type `swap` are not supported for bulk key regeneration.
+                  Bulk Key regeneration will yeild a different output format from single module key regeneration.
     examples:
       - name: Renew the primary key.
         text: az iot hub module-identity renew-key -m {module_name} -d {device_id} -n {iothub_name} --kt primary

azext_iot/operations/hub.py

@@ -539,31 +539,38 @@ def iot_device_key_regenerate(
         renew_key_type += "Key"
 
     modules = []
+    devices = []
     if device_ids[0] == "*":
-        devices = _iot_device_twin_list(target=target, top=None)
-        devices.extend(_iot_device_twin_list(target=target, edge_enabled=True, top=None))
-        device_ids = []
-        for device in devices:
+        device_twins = _iot_device_twin_list(target=target, top=None)
+        device_twins.extend(_iot_device_twin_list(target=target, edge_enabled=True, top=None))
+        for device in device_twins:
             if device["authenticationType"] == DeviceAuthApiType.sas.value:
-                device_ids.append(device["deviceId"])
+                devices.append({"id": device["deviceId"]})
             # non sas devices can have sas modules...
             if include_modules:
                 modules.extend(
                     _iot_key_regenerate_process_modules(
                         target=target, device_id=device["deviceId"], module_ids="*"
                     )
                 )
+    else:
+        devices = [{"id": device_id} for device_id in device_ids]
+
     if modules:
         logger.info(f"Found {len(modules)} modules.")
 
-    # call friendly format
-    devices = [{"id": device_ids[i]} for i in range(len(device_ids))]
-    return _iot_key_regenerate_batch(
+    result = _iot_key_regenerate_batch(
         service_sdk=service_sdk,
         renew_key_type=renew_key_type,
         items=devices + modules,
     )
 
+    # avoid breaking changes by having one device return the device identity
+    if all({len(device_ids) == 1, device_ids[0] != "*", not include_modules}):
+        return _iot_device_show(target, device_ids[0])
+
+    return result
+
 
 def iot_device_get_parent(
     cmd,
@@ -999,12 +1006,7 @@ def iot_device_module_key_regenerate(
             raise InvalidArgumentValueError(
                 "Currently, bulk key swap is not supported."
             )
-        try:
-            module = service_sdk.modules.get_identity(
-                id=device_id, mid=module_ids[0], raw=True
-            ).response.json()
-        except CloudError as e:
-            handle_service_exception(e)
+        module = _iot_device_module_show(target=target, device_id=device_id, module_id=module_ids[0])
         if module["authentication"]["type"] != DeviceAuthApiType.sas.value:
             raise ClientRequestError("Module authentication should be of type sas")
 
@@ -1032,13 +1034,17 @@ def iot_device_module_key_regenerate(
     modules = _iot_key_regenerate_process_modules(
         target, device_id, module_ids
     )
-    return _iot_key_regenerate_batch(
+    result = _iot_key_regenerate_batch(
         service_sdk=service_sdk,
         renew_key_type=renew_key_type,
         items=modules,
         device_id=device_id
     )
 
+    if len(module_ids) == 1 and module_ids[0] != "*":
+        return _iot_device_module_show(target=target, device_id=device_id, module_id=module_ids[0])
+    return result
+
 
 def _iot_key_regenerate_process_modules(
     target,
@@ -1056,7 +1062,7 @@ def _iot_key_regenerate_process_modules(
             if module["authentication"]["type"] == DeviceAuthApiType.sas.value:
                 module_ids.append(module["moduleId"])
 
-    return [{"id": device_id, "moduleId": module_ids[i]} for i in range(len(module_ids))]
+    return [{"id": device_id, "moduleId": module_id} for module_id in module_ids]
 
 
 def _iot_key_regenerate_batch(
@@ -1065,6 +1071,7 @@ def _iot_key_regenerate_batch(
     items,
     device_id=None,
 ):
+    from time import sleep
     overall_result = {
         "policyKey": renew_key_type,
         "errors": [],
@@ -1082,15 +1089,22 @@ def _iot_key_regenerate_batch(
         else:
             batches.append(items[:])
             items = []
-    for i in tqdm(range(len(batches)), desc="Bulk key regeneration is in progress", ascii=' #'):
+    for batch in tqdm(batches, desc="Bulk key regeneration is in progress", ascii=' #'):
         # call
-        try:
-            result = service_sdk.service.bulk_regenerate_device_key_method(
-                policy_key=renew_key_type,
-                devices=batches[i]
-            )
-        except CloudError as e:
-            handle_service_exception(e)
+        result = None
+        tries = 0
+        while tries < 3:
+            try:
+                result = service_sdk.service.bulk_regenerate_device_key_method(
+                    policy_key=renew_key_type,
+                    devices=batch
+                )
+                break
+            except CloudError as e:
+                tries += 1
+                if tries == 3 or e.status_code != 429:
+                    handle_service_exception(e)
+                sleep(20)
         # combine result
         if result.errors:
             overall_result["errors"].extend(result.errors)

azext_iot/tests/digitaltwins/test_dt_model_unit.py

@@ -153,8 +153,7 @@ def post_request_callback_first(request):
 
         def post_request_callback_second(request):
             headers = {"content_type": "application/json"}
-            resp_body = [{"status": "failed"}]
-            return (400, headers, json.dumps(resp_body))
+            return (400, headers, None)
 
         responses.add_callback(
             responses.POST,
@@ -184,12 +183,10 @@ def delete_request_callback(request):
             model = url.split("/")[-1]
             # Ensures that we are deleting the models which were added in the successful batch
             if model in models_added:
-                resp_body = [{"status": "succeeded"}]
                 models_deleted.append(model)
-                return (204, headers, json.dumps(resp_body))
+                return (204, headers, None)
             else:
-                resp_body = [{"status": "Failed - Unexpected model deletion"}]
-                return (400, headers, json.dumps(resp_body))
+                return (400, headers, None)
 
         responses.add_callback(
             responses.DELETE,