fix: az iot hub state export endpoint fix (#705)

* fixes

* passing tests

* history

* make sure tests pass

HISTORY.rst

@@ -15,6 +15,10 @@ Release History
 
 * Removal of `az iot product`.
 
+**IoT Hub updates**
+
+* Fix for `az iot hub state export` to support storage account endpoints with system-assigned identity.
+
 
 0.23.1
 +++++++++++++++

azext_iot/iothub/providers/state.py

@@ -659,10 +659,10 @@ def check_controlplane(self, hub_resource: dict):
                         ep["primaryKey"] = parse_cosmos_db_connection_string(cs_object["connectionString"])["AccountKey"]
                     if cs_object["description"] == "Secondary SQL Connection String" and ep.get("secondaryKey"):
                         ep["secondaryKey"] = parse_cosmos_db_connection_string(cs_object["connectionString"])["AccountKey"]
-            elif isinstance(ep.get("identity"), dict) and ep["identity"]["userAssignedIdentity"] in removed_identities:
+            elif isinstance(ep.get("identity"), dict) and ep["identity"].get("userAssignedIdentity") in removed_identities:
                 logger.warning(
                     usr_msgs.SAVE_ENDPOINT_UAI_RETRIEVE_FAIL_MSG.format(
-                        "Cosmos DB Sql Collection", ep["name"], ep["identity"]["userAssignedIdentity"]
+                        "Cosmos DB Sql Collection", ep["name"], ep["identity"].get("userAssignedIdentity")
                     )
                 )
                 removed_endpoints.append(ep["name"])
@@ -708,10 +708,10 @@ def check_controlplane(self, hub_resource: dict):
                     logger.warning(usr_msgs.SAVE_ENDPOINT_RETRIEVE_FAIL_MSG.format("Event Hub", ep["name"]))
                     removed_endpoints.append(ep["name"])
                     continue
-            elif isinstance(ep.get("identity"), dict) and ep["identity"]["userAssignedIdentity"] in removed_identities:
+            elif isinstance(ep.get("identity"), dict) and ep["identity"].get("userAssignedIdentity") in removed_identities:
                 logger.warning(
                     usr_msgs.SAVE_ENDPOINT_UAI_RETRIEVE_FAIL_MSG.format(
-                        "Event Hub", ep["name"], ep["identity"]["userAssignedIdentity"]
+                        "Event Hub", ep["name"], ep["identity"].get("userAssignedIdentity")
                     )
                 )
                 removed_endpoints.append(ep["name"])
@@ -754,10 +754,10 @@ def check_controlplane(self, hub_resource: dict):
                     logger.warning(usr_msgs.SAVE_ENDPOINT_RETRIEVE_FAIL_MSG.format("Service Bus Queue", ep["name"]))
                     removed_endpoints.append(ep["name"])
                     continue
-            elif isinstance(ep.get("identity"), dict) and ep["identity"]["userAssignedIdentity"] in removed_identities:
+            elif isinstance(ep.get("identity"), dict) and ep["identity"].get("userAssignedIdentity") in removed_identities:
                 logger.warning(
                     usr_msgs.SAVE_ENDPOINT_UAI_RETRIEVE_FAIL_MSG.format(
-                        "Service Bus Queue", ep["name"], ep["identity"]["userAssignedIdentity"]
+                        "Service Bus Queue", ep["name"], ep["identity"].get("userAssignedIdentity")
                     )
                 )
                 removed_endpoints.append(ep["name"])
@@ -802,10 +802,10 @@ def check_controlplane(self, hub_resource: dict):
                     )
                     removed_endpoints.append(ep["name"])
                     continue
-            elif isinstance(ep.get("identity"), dict) and ep["identity"]["userAssignedIdentity"] in removed_identities:
+            elif isinstance(ep.get("identity"), dict) and ep["identity"].get("userAssignedIdentity") in removed_identities:
                 logger.warning(
                     usr_msgs.SAVE_ENDPOINT_UAI_RETRIEVE_FAIL_MSG.format(
-                        "Service Bus Topic", ep["name"], ep["identity"]["userAssignedIdentity"]
+                        "Service Bus Topic", ep["name"], ep["identity"].get("userAssignedIdentity")
                     )
                 )
                 removed_endpoints.append(ep["name"])
@@ -848,10 +848,10 @@ def check_controlplane(self, hub_resource: dict):
                     )
                     removed_endpoints.append(ep["name"])
                     continue
-            elif isinstance(ep.get("identity"), dict) and ep["identity"]["userAssignedIdentity"] in removed_identities:
+            elif isinstance(ep.get("identity"), dict) and ep["identity"].get("userAssignedIdentity") in removed_identities:
                 logger.warning(
                     usr_msgs.SAVE_ENDPOINT_UAI_RETRIEVE_FAIL_MSG.format(
-                        "Storage Container", ep["name"], ep["identity"]["userAssignedIdentity"]
+                        "Storage Container", ep["name"], ep["identity"].get("userAssignedIdentity")
                     )
                 )
                 removed_endpoints.append(ep["name"])
@@ -888,10 +888,10 @@ def check_controlplane(self, hub_resource: dict):
         file_upload = hub_resource["properties"]["storageEndpoints"].get("$default", {})
         if (
             isinstance(file_upload.get("identity"), dict)
-            and file_upload["identity"]["userAssignedIdentity"] in removed_identities
+            and file_upload["identity"].get("userAssignedIdentity") in removed_identities
         ):
             logger.warning(
-                usr_msgs.SAVE_FILE_UPLOAD_UAI_RETRIEVE_FAIL_MSG.format(file_upload["identity"]["userAssignedIdentity"])
+                usr_msgs.SAVE_FILE_UPLOAD_UAI_RETRIEVE_FAIL_MSG.format(file_upload["identity"].get("userAssignedIdentity"))
             )
             file_upload["authenticationType"] = None
             file_upload["connectionString"] = None

azext_iot/tests/iothub/conftest.py

@@ -114,6 +114,7 @@ def setup_hub_controlplane_states(
     provisioned_storage = provisioned_iot_hubs_with_storage_user_module[0]["storage"]
     storage_cstring = provisioned_storage["connectionString"]
     storage_container = provisioned_storage["container"]["name"]
+    storage_endpoint_uri = provisioned_storage["storage"]["primaryEndpoints"]["blob"]
 
     cosmosdb_container_name = provisioned_cosmos_db_module["container"]["name"]
     cosmosdb_database_name = provisioned_cosmos_db_module["database"]["name"]
@@ -176,7 +177,14 @@ def setup_hub_controlplane_states(
 
     cli.invoke(
         f"iot hub message-endpoint create storage-container --en storagecontainer-key -g {hub_rg} -n {hub_name} -c "
-        f"{storage_cstring} --container {storage_container}  -b 350 -w 250 --encoding json")
+        f"{storage_cstring} --container {storage_container}  -b 350 -w 250 --encoding json"
+    )
+
+    cli.invoke(
+        f"iot hub message-endpoint create storage-container --en storagecontainer-{suffix} -g {hub_rg} -n {hub_name} "
+        f"--identity {user_identity_parameter} --endpoint-uri {storage_endpoint_uri} --container {storage_container} "
+        "-b 350 -w 250 --encoding json"
+    )
 
     cli.invoke(
         f"iot hub message-endpoint create cosmosdb-container --en cosmosdb-key -g {hub_rg} -n {hub_name} --container "

azext_iot/tests/iothub/state/test_hub_state_int.py

@@ -259,6 +259,9 @@ def delete_system_endpoints(hub_name, rg):
     cli.invoke(
         f"iot hub routing-endpoint delete --hub-name {hub_name} -g {rg} -n queue-systemid"
     )
+    cli.invoke(
+        f"iot hub routing-endpoint delete --hub-name {hub_name} -g {rg} -n storagecontainer-systemid"
+    )
 
 
 @pytest.mark.hub_infrastructure(count=2, sys_identity=True, user_identity=True, storage=True, desired_tags="abc=def")