From 08b810de67a98d8d21584cf4f3511bf2ce46460b Mon Sep 17 00:00:00 2001
From: Luke Luong <lukeluong@microsoft.com>
Date: Tue, 29 Oct 2024 14:00:30 -0700
Subject: [PATCH] Change performance diagnostics to use Managed Identity

---
 .../azureperformancediagnostics.py            | 23 +++++++++++--------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/microsoft/testsuites/vm_extensions/azureperformancediagnostics.py b/microsoft/testsuites/vm_extensions/azureperformancediagnostics.py
index 5548ab1be4..1525a7a105 100644
--- a/microsoft/testsuites/vm_extensions/azureperformancediagnostics.py
+++ b/microsoft/testsuites/vm_extensions/azureperformancediagnostics.py
@@ -28,7 +28,7 @@
     AzureNodeSchema,
     check_or_create_storage_account,
     get_node_context,
-    get_storage_credential,
+    add_system_assign_identity,
     list_blobs,
 )
 from lisa.sut_orchestrator.azure.features import AzureExtension
@@ -91,12 +91,15 @@ def verify_azure_performance_diagnostics(
             log=log,
         )
 
-        account_credential = get_storage_credential(
-            credential=platform.credential,
-            subscription_id=platform.subscription_id,
-            cloud=platform.cloud,
-            account_name=storage_account_name,
-            resource_group_name=resource_group_name,
+        # Assign system identity to VM
+        # This MSI will be pre-configured with the necessary role assigments to Storage Account from Subscription level
+        # Permissions required: Storage Account Contributor, Storage Blob Data Contributor, Storage Blob Data Contributor
+        add_system_assign_identity(
+            platform=platform,
+            resource_group_name=node_context.resource_group_name,
+            vm_name=node_context.vm_name,
+            location=node_context.location,
+            log=log,
         )
 
         # Run VM Extension
@@ -116,7 +119,7 @@ def verify_azure_performance_diagnostics(
 
         protected_settings = {
             "storageAccountName": storage_account_name,
-            "storageAccountKey": account_credential.get("account_key"),
+            "authenticationType": "SystemManagedIdentity"
         }
 
         extension_result = extension.create_or_update(
@@ -167,11 +170,11 @@ def _is_supported_linux_distro(self, node: Node) -> bool:
             CentOs: [6, 7],
             Oracle: [6, 7],
             Debian: [8, 9, 10, 11],
-            Ubuntu: [14, 16, 18, 20],
+            Ubuntu: [14, 16, 18, 20, 22],
             Suse: [12, 15],
             SLES: [12, 15],
             AlmaLinux: [8],
-            CBLMariner: [2],
+            CBLMariner: [2, 3], #AzureLinux uses CBLMariner class
         }
 
         for distro in supported_major_versions: