mgmt, fix compute MSI sample (Azure#27849)

Author: weidongxu-microsoft

sdk/resourcemanager/azure-resourcemanager-samples/src/main/java/com/azure/resourcemanager/compute/samples/ManageResourceFromMSIEnabledVirtualMachineBelongsToAADGroup.java

@@ -51,21 +51,23 @@ public static boolean runSample(AzureResourceManager azureResourceManager) {
         final String pipName = Utils.randomResourceName(azureResourceManager, "pip1", 15);
         final String userName = "tirekicker";
         final String sshPublicKey = Utils.sshPublicKey();
-        final Region region = Region.US_SOUTH_CENTRAL;
+        final Region region = Region.US_EAST;
 
         final String installScript = "https://raw.githubusercontent.com/Azure/azure-sdk-for-java/main/sdk/resourcemanager/azure-resourcemanager-samples/src/main/resources/create_resources_with_msi.sh";
         String installCommand = "bash create_resources_with_msi.sh {stgName} {rgName} {location}";
         List<String> fileUris = new ArrayList<>();
         fileUris.add(installScript);
 
+        ActiveDirectoryGroup activeDirectoryGroup = null;
+
         try {
 
             //=============================================================
             // Create a AAD security group
 
             System.out.println("Creating a AAD security group");
 
-            ActiveDirectoryGroup activeDirectoryGroup = azureResourceManager.accessManagement()
+            activeDirectoryGroup = azureResourceManager.accessManagement()
                     .activeDirectoryGroups()
                     .define(groupName)
                         .withEmailAlias(groupName)
@@ -127,11 +129,11 @@ public static boolean runSample(AzureResourceManager azureResourceManager) {
 
             System.out.println("Added virtual machine MSI service principal to the AAD group");
 
-            System.out.println("Waiting 15 minutes to MSI extension in the VM to refresh the token");
+            System.out.println("Waiting 10 minutes to MSI extension in the VM to refresh the token");
 
             ResourceManagerUtils.sleep(Duration.ofMinutes(10));
 
-            // Prepare custom script t install az cli that uses MSI to create a storage account
+            // Prepare custom script to install az cli that uses MSI to create a storage account
             //
             final String stgName = Utils.randomResourceName(azureResourceManager, "st44", 15);
             installCommand = installCommand.replace("{stgName}", stgName)
@@ -145,10 +147,10 @@ public static boolean runSample(AzureResourceManager azureResourceManager) {
 
             virtualMachine
                     .update()
-                        .defineNewExtension("CustomScriptForLinux")
-                            .withPublisher("Microsoft.OSTCExtensions")
-                            .withType("CustomScriptForLinux")
-                            .withVersion("1.4")
+                        .defineNewExtension("az-create-storage-account")
+                            .withPublisher("Microsoft.Azure.Extensions")
+                            .withType("CustomScript")
+                            .withVersion("2.1")
                             .withMinorVersionAutoUpgrade()
                             .withPublicSetting("fileUris", fileUris)
                             .withPublicSetting("commandToExecute", installCommand)
@@ -164,6 +166,15 @@ public static boolean runSample(AzureResourceManager azureResourceManager) {
             Utils.print(storageAccount);
             return true;
         } finally {
+            try {
+                if (activeDirectoryGroup != null) {
+                    System.out.println("Deleting AAD Group: " + activeDirectoryGroup.name());
+                    azureResourceManager.accessManagement().activeDirectoryGroups()
+                        .deleteById(activeDirectoryGroup.id());
+                }
+            } catch (Exception e) {
+                e.printStackTrace();
+            }
             try {
                 System.out.println("Deleting Resource Group: " + rgName);
                 azureResourceManager.resourceGroups().beginDeleteByName(rgName);

sdk/resourcemanager/azure-resourcemanager-samples/src/main/java/com/azure/resourcemanager/compute/samples/ManageStorageFromMSIEnabledVirtualMachine.java

@@ -42,7 +42,7 @@ public static boolean runSample(AzureResourceManager azureResourceManager) {
         final String pipName = Utils.randomResourceName(azureResourceManager, "pip1", 15);
         final String userName = "tirekicker";
         final String sshPublicKey = Utils.sshPublicKey();
-        final Region region = Region.US_WEST_CENTRAL;
+        final Region region = Region.US_EAST;
 
         final String installScript = "https://raw.githubusercontent.com/Azure/azure-sdk-for-java/main/sdk/resourcemanager/azure-resourcemanager-samples/src/main/resources/create_resources_with_msi.sh";
         String installCommand = "bash create_resources_with_msi.sh {stgName} {rgName} {location}";
@@ -74,7 +74,7 @@ public static boolean runSample(AzureResourceManager azureResourceManager) {
             System.out.println("Created virtual machine with MSI enabled");
             Utils.print(virtualMachine);
 
-            // Prepare custom script t install az cli that uses MSI to create a storage account
+            // Prepare custom script to install az cli that uses MSI to create a storage account
             //
             final String stgName = Utils.randomResourceName(azureResourceManager, "st44", 15);
             installCommand = installCommand.replace("{stgName}", stgName)
@@ -88,10 +88,10 @@ public static boolean runSample(AzureResourceManager azureResourceManager) {
 
             virtualMachine
                     .update()
-                        .defineNewExtension("CustomScriptForLinux")
-                            .withPublisher("Microsoft.OSTCExtensions")
-                            .withType("CustomScriptForLinux")
-                            .withVersion("1.4")
+                        .defineNewExtension("az-create-storage-account")
+                            .withPublisher("Microsoft.Azure.Extensions")
+                            .withType("CustomScript")
+                            .withVersion("2.1")
                             .withMinorVersionAutoUpgrade()
                             .withPublicSetting("fileUris", fileUris)
                             .withPublicSetting("commandToExecute", installCommand)

sdk/resourcemanager/azure-resourcemanager-samples/src/main/resources/create_resources_with_msi.sh

@@ -1,8 +1,6 @@
 #!/bin/bash
 
-/usr/bin/yes | sudo apt-get update
-/usr/bin/yes | sudo apt install python-pip
-/usr/bin/yes | sudo pip install --upgrade pip
-sudo pip install azure-cli
-az login --msi
-az storage account create -n $1 -g $2 -l $3 --sku Premium_LRS
+curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash
+
+az login --identity
+az storage account create -n $1 -g $2 -l $3 --sku Premium_LRS

sdk/resourcemanager/azure-resourcemanager-samples/src/test/java/com/azure/resourcemanager/samples/ComputeSampleTests.java

@@ -150,13 +150,11 @@ public void testManageManagedDisks() {
     }
 
     @Test
-    @Disabled("Skipping for now - looks like a service side issue")
     public void testManageStorageFromMSIEnabledVirtualMachine() {
         Assertions.assertTrue(ManageStorageFromMSIEnabledVirtualMachine.runSample(azureResourceManager));
     }
 
     @Test
-    @Disabled("Skipping for now - looks like a service side issue")
     public void testManageResourceFromMSIEnabledVirtualMachineBelongsToAADGroup() {
         Assertions.assertTrue(ManageResourceFromMSIEnabledVirtualMachineBelongsToAADGroup.runSample(azureResourceManager));
     }