[PSCmdAssistant] Disk Policy for Managed Disks

src/Compute/Compute.Test/ScenarioTests/DiskRPTests.cs

@@ -199,5 +199,19 @@ public void TestDiskGrantAccessGetSASWithTL()
             TestRunner.RunTestScript("Test-DiskGrantAccessGetSASWithTL");
         }
 
-    }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void testgennewazdiskupdateconfig()
+        {
+            TestRunner.RunTestScript("TestGen-newazdiskupdateconfig");
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void testgennewazdisk()
+        {
+            TestRunner.RunTestScript("TestGen-newazdisk");
+        }
+            }
 }

src/Compute/Compute.Test/ScenarioTests/DiskRPTests.ps1

@@ -1863,4 +1863,82 @@ function Test-DiskGrantAccessGetSASWithTL
 		# Cleanup
 		Clean-ResourceGroup $rgname;
 	}
-}
+}
+function TestGen-newazdiskupdateconfig
+{
+    $rgname = Get-ComputeTestResourceName;
+    $loc = Get-Location;
+
+    try
+    {
+        New-AzResourceGroup -Name $rgname -Location $loc -Force;
+
+        # Create a new disk with initial configuration
+        $diskConfig = New-AzDiskConfig -Location $loc -SkuName 'Premium_LRS' -CreateOption 'Empty' -DiskSizeGB 2;
+        $diskname = "disk" + $rgname;
+        $diskPr = New-AzDisk -ResourceGroupName $rgname -DiskName $diskname -Disk $diskConfig;
+
+        # Verify initial disk creation
+        $disk = Get-AzDisk -ResourceGroupName $rgname -DiskName $diskname;
+        Assert-AreEqual $disk.Sku.Name "Premium_LRS";
+
+        # Update the disk with new AvailabilityPolicy parameter
+        $diskupdateconfig = New-AzDiskUpdateConfig -AvailabilityPolicy 'AutomaticReattach';
+        Update-AzDisk -ResourceGroupName $rgname -DiskName $diskname -DiskUpdate $diskupdateconfig;
+
+        # Verify the disk update
+        $updatedDisk = Get-AzDisk -ResourceGroupName $rgname -DiskName $diskname;
+        Assert-AreEqual $updatedDisk.AvailabilityPolicy "AutomaticReattach";
+
+        # Update the disk to clear the AvailabilityPolicy
+        $diskupdateconfig = New-AzDiskUpdateConfig -AvailabilityPolicy 'None';
+        Update-AzDisk -ResourceGroupName $rgname -DiskName $diskname -DiskUpdate $diskupdateconfig;
+
+        # Verify the disk update
+        $updatedDisk = Get-AzDisk -ResourceGroupName $rgname -DiskName $diskname;
+        Assert-AreEqual $updatedDisk.AvailabilityPolicy $null;
+    }
+    finally
+    {
+        # Clean up resources
+        Remove-AzResourceGroup -Name $rgname -Force -ErrorAction SilentlyContinue;
+    }
+}
+
+function TestGen-newazdisk
+{
+    $rgname = Get-ComputeTestResourceName;
+    $loc = Get-Location;
+
+    try
+    {
+        New-AzResourceGroup -Name $rgname -Location $loc -Force;
+
+        # Setup
+        $diskName = 'disk' + $rgname;
+        $diskAccountType = 'Premium_LRS';
+        $createOption = 'Empty';
+        $diskSize = 32;
+        $availabilityPolicy1 = 'AutomaticReattach';
+        $availabilityPolicy2 = 'None';
+
+        # Create Disk with AvailabilityPolicy
+        $diskConfig = New-AzDiskConfig -Location $loc -AccountType $diskAccountType -CreateOption $createOption -DiskSizeGB $diskSize -AvailabilityPolicy $availabilityPolicy1;
+        New-AzDisk -ResourceGroupName $rgname -DiskName $diskName -Disk $diskConfig;
+        $disk = Get-AzDisk -ResourceGroupName $rgname -DiskName $diskName;
+        Assert-AreEqual $disk.AvailabilityPolicy $availabilityPolicy1;
+
+        # Update Disk with new AvailabilityPolicy
+        $diskUpdateConfig = New-AzDiskUpdateConfig -AvailabilityPolicy $availabilityPolicy2;
+        Update-AzDisk -ResourceGroupName $rgname -DiskName $diskName -DiskUpdate $diskUpdateConfig;
+        $diskUpdated = Get-AzDisk -ResourceGroupName $rgname -DiskName $diskName;
+        Assert-AreEqual $diskUpdated.AvailabilityPolicy $availabilityPolicy2;
+
+        # Clean up
+        Remove-AzDisk -ResourceGroupName $rgname -DiskName $diskName -Force;
+    }
+    finally
+    {
+        Remove-AzResourceGroup -Name $rgname -Force -ErrorAction SilentlyContinue;
+    }
+}

src/Compute/Compute.Test/ScenarioTests/VirtualMachineTests.cs

@@ -675,5 +675,19 @@ public void TestVMSetAzOSCredentialNullRef()
         {
             TestRunner.RunTestScript("Test-VMSetAzOSCredentialNullRef");
         }
-    }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void testgensetazvm()
+        {
+            TestRunner.RunTestScript("TestGen-setazvm");
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void testgensetazvmdatadisk()
+        {
+            TestRunner.RunTestScript("TestGen-setazvmdatadisk");
+        }
+            }
 }

src/Compute/Compute.Test/ScenarioTests/VirtualMachineTests.ps1

@@ -7727,3 +7727,135 @@ function Test-VMwithSSHKeyEd25519
         Clean-ResourceGroup $rgname;
     }
 }
+
+function TestGen-setazvm
+{
+    # Setup
+    $rgname = Get-ComputeTestResourceName;
+    $loc = Get-ComputeVMLocation;
+
+    try
+    {
+        New-AzResourceGroup -Name $rgname -Location $loc -Force;
+
+        # VM Profile & Hardware
+        $vmname = 'vm' + $rgname;
+        $diskName = $vmname + "-datadisk";
+        $diskSizeGB = 1024;
+        $availabilityPolicy = "AutomaticReattach";
+
+        # Create a managed disk
+        $diskConfig = New-AzDiskConfig -Location $loc -CreateOption Empty -DiskSizeGB $diskSizeGB;
+        $disk = New-AzDisk -ResourceGroupName $rgname -DiskName $diskName -Disk $diskConfig;
+
+        # Update the disk with AvailabilityPolicy
+        $diskUpdateConfig = New-AzDiskUpdateConfig -AvailabilityPolicy $availabilityPolicy;
+        Update-AzDisk -ResourceGroupName $rgname -DiskName $diskName -DiskUpdate $diskUpdateConfig;
+
+        # Attach the disk to a VM
+        $vm = Get-AzVM -ResourceGroupName $rgname -Name $vmname;
+        $vm = Add-AzVMDataDisk -VM $vm -Name $diskName -CreateOption Attach -ManagedDiskId $disk.Id;
+        Set-AzVM -ResourceGroupName $rgname -VM $vm;
+
+        # Validate the AvailabilityPolicy on the disk
+        $updatedDisk = Get-AzDisk -ResourceGroupName $rgname -DiskName $diskName;
+        Assert-AreEqual $updatedDisk.AvailabilityPolicy $availabilityPolicy;
+
+        # Test clearing the AvailabilityPolicy
+        $diskUpdateConfigClear = New-AzDiskUpdateConfig -AvailabilityPolicy "None";
+        Update-AzDisk -ResourceGroupName $rgname -DiskName $diskName -DiskUpdate $diskUpdateConfigClear;
+
+        # Validate the AvailabilityPolicy is cleared
+        $clearedDisk = Get-AzDisk -ResourceGroupName $rgname -DiskName $diskName;
+        Assert-AreEqual $clearedDisk.AvailabilityPolicy "None";
+    }
+    finally
+    {
+        # Cleanup
+        Remove-AzResourceGroup -Name $rgname -Force -ErrorAction SilentlyContinue;
+    }
+}
+
+function TestGen-setazvmdatadisk
+{
+    # Setup
+    $rgname = Get-ComputeTestResourceName
+
+    try
+    {
+        # Common
+        $loc = Get-ComputeVMLocation;
+        New-AzResourceGroup -Name $rgname -Location $loc -Force;
+
+        # VM Profile & Hardware
+        $vmsize = "Standard_DS2_v2";
+        $vmname = 'vm' + $rgname;
+        $p = New-AzVMConfig -VMName $vmname -VMSize $vmsize;
+        Assert-AreEqual $p.HardwareProfile.VmSize $vmsize;
+
+        # NRP
+        $subnet = New-AzVirtualNetworkSubnetConfig -Name ('subnet' + $rgname) -AddressPrefix "10.0.0.0/24";
+        $vnet = New-AzVirtualNetwork -Force -Name ('vnet' + $rgname) -ResourceGroupName $rgname -Location $loc -AddressPrefix "10.0.0.0/16" -Subnet $subnet;
+        $vnet = Get-AzVirtualNetwork -Name ('vnet' + $rgname) -ResourceGroupName $rgname;
+        $subnetId = $vnet.Subnets[0].Id;
+        $pubip = New-AzPublicIpAddress -Force -Name ('pubip' + $rgname) -ResourceGroupName $rgname -Location $loc -AllocationMethod Static -DomainNameLabel ('pubip' + $rgname);
+        $pubip = Get-AzPublicIpAddress -Name ('pubip' + $rgname) -ResourceGroupName $rgname;
+        $pubipId = $pubip.Id;
+        $nic = New-AzNetworkInterface -Force -Name ('nic' + $rgname) -ResourceGroupName $rgname -Location $loc -SubnetId $subnetId -PublicIpAddressId $pubip.Id;
+        $nic = Get-AzNetworkInterface -Name ('nic' + $rgname) -ResourceGroupName $rgname;
+        $nicId = $nic.Id;
+
+        $p = Add-AzVMNetworkInterface -VM $p -Id $nicId;
+        Assert-AreEqual $p.NetworkProfile.NetworkInterfaces.Count 1;
+        Assert-AreEqual $p.NetworkProfile.NetworkInterfaces[0].Id $nicId;
+
+        # Storage Account (SA)
+        $stoname = 'sto' + $rgname;
+        $stotype = 'Standard_GRS';
+        New-AzStorageAccount -ResourceGroupName $rgname -Name $stoname -Location $loc -Type $stotype;
+        $stoaccount = Get-AzStorageAccount -ResourceGroupName $rgname -Name $stoname;
+
+        $osDiskName = 'osDisk';
+        $osDiskCaching = 'ReadWrite';
+        $osDiskVhdUri = "https://$stoname.blob.core.windows.net/test/os.vhd";
+        $dataDiskVhdUri1 = "https://$stoname.blob.core.windows.net/test/data1.vhd";
+
+        $p = Set-AzVMOSDisk -VM $p -Name $osDiskName -VhdUri $osDiskVhdUri -Caching $osDiskCaching -CreateOption FromImage;
+        $p = Add-AzVMDataDisk -VM $p -Name 'testDataDisk1' -Caching 'ReadOnly' -DiskSizeInGB 10 -Lun 1 -VhdUri $dataDiskVhdUri1 -CreateOption Empty -AvailabilityPolicy 'AutomaticReattach';
+
+        Assert-AreEqual $p.StorageProfile.OSDisk.Caching $osDiskCaching;
+        Assert-AreEqual $p.StorageProfile.OSDisk.Name $osDiskName;
+        Assert-AreEqual $p.StorageProfile.OSDisk.Vhd.Uri $osDiskVhdUri;
+        Assert-AreEqual $p.StorageProfile.DataDisks.Count 1;
+        Assert-AreEqual $p.StorageProfile.DataDisks[0].Caching 'ReadOnly';
+        Assert-AreEqual $p.StorageProfile.DataDisks[0].DiskSizeGB 10;
+        Assert-AreEqual $p.StorageProfile.DataDisks[0].Lun 1;
+        Assert-AreEqual $p.StorageProfile.DataDisks[0].Vhd.Uri $dataDiskVhdUri1;
+        Assert-AreEqual $p.StorageProfile.DataDisks[0].AvailabilityPolicy 'AutomaticReattach';
+
+        # Virtual Machine
+        New-AzVM -ResourceGroupName $rgname -Location $loc -VM $p;
+
+        # Get VM
+        $vm1 = Get-AzVM -Name $vmname -ResourceGroupName $rgname;
+        Assert-AreEqual $vm1.Name $vmname;
+        Assert-AreEqual $vm1.StorageProfile.DataDisks[0].AvailabilityPolicy 'AutomaticReattach';
+
+        # Update Data Disk AvailabilityPolicy
+        $vm1 = Set-AzVMDataDisk -VM $vm1 -Name 'testDataDisk1' -AvailabilityPolicy 'None';
+        Update-AzVM -ResourceGroupName $rgname -VM $vm1;
+
+        $vm2 = Get-AzVM -Name $vmname -ResourceGroupName $rgname;
+        Assert-AreEqual $vm2.StorageProfile.DataDisks[0].AvailabilityPolicy 'None';
+
+        # Remove All VMs
+        Get-AzVM -ResourceGroupName $rgname | Remove-AzVM -ResourceGroupName $rgname -Force;
+        $vms = Get-AzVM -ResourceGroupName $rgname;
+        Assert-AreEqual $vms $null;
+    }
+    finally
+    {
+        # Cleanup
+        Remove-AzResourceGroup -Name $rgname -Force;
+    }
+}

src/Compute/Compute/ChangeLog.md

@@ -20,6 +20,12 @@
 
 -->
 ## Upcoming Release
+* Added new optional parameter `-AvailabilityPolicy` to `Set-AzVMDataDisk` and `New-AzDiskUpdateConfig` cmdlets.
+    - The `-AvailabilityPolicy` parameter accepts values 'AutomaticReattach' and 'None', allowing users to set or clear the availability policy on disks.
+
+* Upgraded nuget package to signed package.
+* Upgraded Azure.Core to 1.44.1.
+* Compute gallery related cmdlets will now use 2024-03-03 GalleryRP API calls. 
 * Upgraded nuget package to signed package.
 * Upgraded Azure.Core to 1.44.1.
 * Compute gallery related cmdlets will now use 2024-03-03 GalleryRP API calls. 

src/Compute/Compute/Generated/Disk/Config/NewAzureRmDiskUpdateConfigCommand.cs

@@ -1,4 +1,4 @@
-//
+ //
 // Copyright (c) Microsoft and contributors.  All rights reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -179,6 +179,12 @@ public partial class NewAzureRmDiskUpdateConfigCommand : Microsoft.Azure.Command
         [PSArgumentCompleter("X64", "Arm64")]
         public string Architecture { get; set; }
 
+        [Parameter(
+            Mandatory = false,
+            ValueFromPipelineByPropertyName = true,
+            HelpMessage = "2 values - 'AutomaticReattach': Will set/update property on disk and 'None': Will clear property on disk")]
+        [PSArgumentCompleter("AutomaticReattach", "None")]
+        public string AvailabilityPolicy { get; set; }
 
         protected override void ProcessRecord()
         {
@@ -322,7 +328,8 @@ private void Run()
                 SupportsHibernation = this.IsParameterBound(c => c.SupportsHibernation) ? SupportsHibernation : null,
                 SupportedCapabilities = vSupportedCapabilities,
                 PublicNetworkAccess = this.IsParameterBound(c => c.PublicNetworkAccess) ? PublicNetworkAccess : null,
-                DataAccessAuthMode = this.IsParameterBound(c => c.DataAccessAuthMode) ? DataAccessAuthMode : null
+                DataAccessAuthMode = this.IsParameterBound(c => c.DataAccessAuthMode) ? DataAccessAuthMode : null,
+                AvailabilityPolicy = this.IsParameterBound(c => c.AvailabilityPolicy) ? this.AvailabilityPolicy : null
             };
 
             WriteObject(vDiskUpdate);

src/Compute/Compute/Generated/Models/PSDiskUpdate.cs

@@ -57,5 +57,7 @@ public partial class PSDiskUpdate
         //     Gets or sets possible values include: 'Enabled', 'Disabled'
         public string PublicNetworkAccess { get; set; }
         public string DataAccessAuthMode { get; set; }
+
+        public string AvailabilityPolicy { get; set; }
     }
-}
+}

src/Compute/Compute/VirtualMachine/Config/SetAzureVMDataDiskCommand.cs

@@ -1,4 +1,4 @@
-// ----------------------------------------------------------------------------------
+// ----------------------------------------------------------------------------------
 //
 // Copyright Microsoft Corporation
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -92,6 +92,12 @@ public class SetAzureVMDataDiskCommand : Microsoft.Azure.Commands.ResourceManage
             ValueFromPipelineByPropertyName = false)]
         public SwitchParameter WriteAccelerator { get; set; }
 
+        [Parameter(
+            Mandatory = false,
+            HelpMessage = "2 values - 'AutomaticReattach': Will set/update property on disk and 'None': Will clear property on disk")]
+        [ValidateSet("AutomaticReattach", "None", IgnoreCase = true)]
+        public string AvailabilityPolicy { get; set; }
+
         public override void ExecuteCmdlet()
         {
             var storageProfile = this.VM.StorageProfile;
@@ -152,6 +158,11 @@ public override void ExecuteCmdlet()
                     }
                 }
 
+                if (this.IsParameterBound(c => c.AvailabilityPolicy))
+                {
+                    dataDisk.ManagedDisk.AvailabilityPolicy = this.AvailabilityPolicy;
+                }
+
                 dataDisk.WriteAcceleratorEnabled = this.WriteAccelerator.IsPresent;
             }