From 15aefb2d259cb57b2623020bc2ad5e484bda7344 Mon Sep 17 00:00:00 2001
From: "Hong Li(MSFT)" <74638143+v-hongli1@users.noreply.github.com>
Date: Wed, 20 Nov 2024 12:20:58 +0800
Subject: [PATCH] Fix for introduce WithCapacityReservation DefinitionStage
 into the VirtualMachine model (#42914)

* Add support for capacity reservation.
---
 .../CHANGELOG.md                              |  6 +-
 .../implementation/VirtualMachineImpl.java    | 23 +++++
 .../compute/models/VirtualMachine.java        | 35 ++++++-
 .../VirtualMachineOperationsTests.java        | 94 +++++++++++++++++++
 4 files changed, 151 insertions(+), 7 deletions(-)

diff --git a/sdk/resourcemanager/azure-resourcemanager-compute/CHANGELOG.md b/sdk/resourcemanager/azure-resourcemanager-compute/CHANGELOG.md
index 5bb186b5b4b5..faa9996bc411 100644
--- a/sdk/resourcemanager/azure-resourcemanager-compute/CHANGELOG.md
+++ b/sdk/resourcemanager/azure-resourcemanager-compute/CHANGELOG.md
@@ -4,11 +4,7 @@
 
 ### Features Added
 
-### Breaking Changes
-
-### Bugs Fixed
-
-### Other Changes
+- Supported `capacityReservationGroupId` and `withCapacityReservationGroup` methods for `VirtualMachine`.
 
 ## 2.44.0 (2024-10-25)
 
diff --git a/sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/implementation/VirtualMachineImpl.java b/sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/implementation/VirtualMachineImpl.java
index b233fc4178e3..c8c9d3270a8d 100644
--- a/sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/implementation/VirtualMachineImpl.java
+++ b/sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/implementation/VirtualMachineImpl.java
@@ -27,6 +27,7 @@
 import com.azure.resourcemanager.compute.models.BillingProfile;
 import com.azure.resourcemanager.compute.models.BootDiagnostics;
 import com.azure.resourcemanager.compute.models.CachingTypes;
+import com.azure.resourcemanager.compute.models.CapacityReservationProfile;
 import com.azure.resourcemanager.compute.models.DataDisk;
 import com.azure.resourcemanager.compute.models.DeleteOptions;
 import com.azure.resourcemanager.compute.models.DiagnosticsProfile;
@@ -1993,6 +1994,15 @@ public String userData() {
         return this.innerModel().userData();
     }
 
+    @Override
+    public String capacityReservationGroupId() {
+        if (this.innerModel().capacityReservation() != null
+            && this.innerModel().capacityReservation().capacityReservationGroup() != null) {
+            return this.innerModel().capacityReservation().capacityReservationGroup().id();
+        }
+        return null;
+    }
+
     // CreateUpdateTaskGroup.ResourceCreator.beforeGroupCreateOrUpdate implementation
     @Override
     public void beforeGroupCreateOrUpdate() {
@@ -2728,6 +2738,8 @@ private void copyInnerToUpdateParameter(VirtualMachineUpdateInner updateParamete
         updateParameter.withPriority(this.innerModel().priority());
         updateParameter.withEvictionPolicy(this.innerModel().evictionPolicy());
         updateParameter.withUserData(this.innerModel().userData());
+        updateParameter.withCapacityReservation(this.innerModel().capacityReservation());
+        updateParameter.withApplicationProfile(this.innerModel().applicationProfile());
     }
 
     RoleAssignmentHelper.IdProvider idProvider() {
@@ -2866,6 +2878,17 @@ public VirtualMachineImpl withoutEncryptionAtHost() {
         return this;
     }
 
+    @Override
+    public VirtualMachineImpl withCapacityReservationGroup(String capacityReservationGroupId) {
+        if (this.innerModel().capacityReservation() == null) {
+            this.innerModel().withCapacityReservation(new CapacityReservationProfile());
+        }
+        this.innerModel()
+            .capacityReservation()
+            .withCapacityReservationGroup(new SubResource().withId(capacityReservationGroupId));
+        return this;
+    }
+
     /** Class to manage Data disk collection. */
     private class ManagedDataDiskCollection {
         private final Map<String, DataDisk> newDisksToAttach = new HashMap<>();
diff --git a/sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/models/VirtualMachine.java b/sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/models/VirtualMachine.java
index fb1f87fe67b9..eb7d6210c3ce 100644
--- a/sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/models/VirtualMachine.java
+++ b/sdk/resourcemanager/azure-resourcemanager-compute/src/main/java/com/azure/resourcemanager/compute/models/VirtualMachine.java
@@ -455,6 +455,13 @@ Mono<RunCommandResult> runShellScriptAsync(List<String> scriptLines,
      */
     String userData();
 
+    /**
+     * Gets resource id of capacity reservation group for the virtual machine.
+     *
+     * @return the resource id of capacity reservation group for the virtual machine.
+     */
+    String capacityReservationGroupId();
+
     // Setters
     //
 
@@ -1952,6 +1959,17 @@ interface WithUserData {
             WithCreate withUserData(String base64EncodedUserData);
         }
 
+        /** The stage of a virtual machine definition allowing to specify capacity reservation configurations. */
+        interface WithCapacityReservation {
+            /**
+             * Specifies resource id of capacity reservation group for the virtual machine.
+             *
+             * @param capacityReservationGroupId the resource id of capacity reservation group
+             * @return the next stage of the definition
+             */
+            WithCreate withCapacityReservationGroup(String capacityReservationGroupId);
+        }
+
         /**
          * The stage of the definition which contains all the minimum required inputs for the resource to be created,
          * but also allows for any other optional settings to be specified.
@@ -1965,7 +1983,8 @@ interface WithCreate extends Creatable<VirtualMachine>, Resource.DefinitionWithT
             DefinitionStages.WithUserAssignedManagedServiceIdentity, DefinitionStages.WithLicenseType,
             DefinitionStages.WithAdditionalCapacities, DefinitionStages.WithNetworkInterfaceDeleteOptions,
             DefinitionStages.WithEphemeralOSDisk, DefinitionStages.WithScaleSet, DefinitionStages.WithSecurityTypes,
-            DefinitionStages.WithSecurityProfile, DefinitionStages.WithUserData {
+            DefinitionStages.WithSecurityProfile, DefinitionStages.WithUserData,
+            DefinitionStages.WithCapacityReservation {
 
             /**
              * Begins creating the virtual machine resource.
@@ -2617,6 +2636,17 @@ interface WithUserData {
              */
             Update withUserData(String base64EncodedUserData);
         }
+
+        /** The stage of the virtual machine update allowing to capacity reservation configurations. */
+        interface WithCapacityReservation {
+            /**
+             * Specifies resource id of capacity reservation group for the virtual machine.
+             *
+             * @param capacityReservationGroupId the resource id of capacity reservation group
+             * @return the next stage of the update
+             */
+            Update withCapacityReservationGroup(String capacityReservationGroupId);
+        }
     }
 
     /** The template for an update operation, containing all the settings that can be modified. */
@@ -2626,7 +2656,8 @@ interface Update extends Appliable<VirtualMachine>, Resource.UpdateWithTags<Upda
         UpdateStages.WithBillingProfile, UpdateStages.WithSystemAssignedManagedServiceIdentity,
         UpdateStages.WithUserAssignedManagedServiceIdentity, UpdateStages.WithLicenseType,
         UpdateStages.WithAdditionalCapacities, UpdateStages.WithOSDisk, UpdateStages.WithSecurityFeatures,
-        UpdateStages.WithDeleteOptions, UpdateStages.WithSecurityProfile, UpdateStages.WithUserData {
+        UpdateStages.WithDeleteOptions, UpdateStages.WithSecurityProfile, UpdateStages.WithUserData,
+        UpdateStages.WithCapacityReservation {
         /**
          * Specifies the encryption settings for the OS Disk.
          *
diff --git a/sdk/resourcemanager/azure-resourcemanager-compute/src/test/java/com/azure/resourcemanager/compute/VirtualMachineOperationsTests.java b/sdk/resourcemanager/azure-resourcemanager-compute/src/test/java/com/azure/resourcemanager/compute/VirtualMachineOperationsTests.java
index 2f2da688c8dc..0d48668bcaff 100644
--- a/sdk/resourcemanager/azure-resourcemanager-compute/src/test/java/com/azure/resourcemanager/compute/VirtualMachineOperationsTests.java
+++ b/sdk/resourcemanager/azure-resourcemanager-compute/src/test/java/com/azure/resourcemanager/compute/VirtualMachineOperationsTests.java
@@ -17,6 +17,8 @@
 import com.azure.core.util.logging.LogLevel;
 import com.azure.core.util.polling.LongRunningOperationStatus;
 import com.azure.core.util.polling.PollResponse;
+import com.azure.resourcemanager.compute.fluent.models.CapacityReservationGroupInner;
+import com.azure.resourcemanager.compute.fluent.models.CapacityReservationInner;
 import com.azure.resourcemanager.compute.fluent.models.VirtualMachineInner;
 import com.azure.resourcemanager.compute.models.ApiErrorException;
 import com.azure.resourcemanager.compute.models.AvailabilitySet;
@@ -38,6 +40,7 @@
 import com.azure.resourcemanager.compute.models.RunCommandInputParameter;
 import com.azure.resourcemanager.compute.models.RunCommandResult;
 import com.azure.resourcemanager.compute.models.SecurityTypes;
+import com.azure.resourcemanager.compute.models.Sku;
 import com.azure.resourcemanager.compute.models.UpgradeMode;
 import com.azure.resourcemanager.compute.models.VirtualMachine;
 import com.azure.resourcemanager.compute.models.VirtualMachineDiskOptions;
@@ -77,6 +80,7 @@
 
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -114,6 +118,96 @@ protected void cleanUpResources() {
         }
     }
 
+    @Test
+    @Disabled("No found any available vm size and zone(s) (if zonal) for capacity reservation.")
+    public void canCreateAndUpdateVirtualMachineWithCapacityReservation() {
+        final String crgNameForCreate = generateRandomResourceName("crg", 15);
+        final String crgNameForUpdate = generateRandomResourceName("crg", 15);
+
+        // Create resource group
+        resourceManager.resourceGroups().define(rgName).withRegion(region).create();
+
+        // Create a capacity reservation group for create virtual machine
+        CapacityReservationGroupInner crgForCreate = computeManager.serviceClient()
+            .getCapacityReservationGroups()
+            .createOrUpdate(rgName, crgNameForCreate,
+                new CapacityReservationGroupInner().withLocation(region.name()).withZones(Arrays.asList("1")));
+
+        computeManager.serviceClient()
+            .getCapacityReservations()
+            .createOrUpdate(rgName, crgNameForCreate, generateRandomResourceName("cr", 15),
+                new CapacityReservationInner().withLocation(region.name())
+                    .withZones(Arrays.asList("1"))
+                    .withSku(new Sku().withName("Standard_DS1_v2").withCapacity(4L)));
+
+        // Create another capacity reservation group for update virtual machine
+        CapacityReservationGroupInner crgForUpdate = computeManager.serviceClient()
+            .getCapacityReservationGroups()
+            .createOrUpdate(rgName, crgNameForUpdate,
+                new CapacityReservationGroupInner().withLocation(region.name()).withZones(Arrays.asList("1")));
+
+        // Create network for virtual machine
+        Network network = computeManager.networkManager()
+            .networks()
+            .define(generateRandomResourceName("vnet", 15))
+            .withRegion(region)
+            .withNewResourceGroup(rgName)
+            .withAddressSpace("10.0.0.0/28")
+            .create();
+
+        computeManager.serviceClient()
+            .getCapacityReservations()
+            .createOrUpdate(rgName, crgNameForUpdate, generateRandomResourceName("cr", 15),
+                new CapacityReservationInner().withLocation(region.name())
+                    .withZones(Arrays.asList("1"))
+                    .withSku(new Sku().withName("Standard_DS1_v2").withCapacity(4L)));
+
+        // Create virtual machine without any capacity reservations
+        VirtualMachine vm = computeManager.virtualMachines()
+            .define(generateRandomResourceName("vm", 15))
+            .withRegion(region)
+            .withNewResourceGroup(rgName)
+            .withExistingPrimaryNetwork(network)
+            .withSubnet(generateRandomResourceName("subnet", 15))
+            .withPrimaryPrivateIPAddressDynamic()
+            .withoutPrimaryPublicIPAddress()
+            .withPopularWindowsImage(KnownWindowsVirtualMachineImage.WINDOWS_SERVER_2019_DATACENTER_GEN2)
+            .withAdminUsername("Foo12")
+            .withAdminPassword(password())
+            .withNewDataDisk(127)
+            .withSize(VirtualMachineSizeTypes.STANDARD_DS1_V2)
+            .create();
+
+        // Update virtual machine with capacity reservation group
+        vm.update().withCapacityReservationGroup(crgForUpdate.id()).apply();
+        Assertions.assertEquals(crgForUpdate.id(), vm.capacityReservationGroupId());
+        computeManager.virtualMachines().deleteById(vm.id());
+
+        // Create virtual machine with capacity reservation group
+        vm = computeManager.virtualMachines()
+            .define(generateRandomResourceName("vm", 15))
+            .withRegion(region)
+            .withNewResourceGroup(rgName)
+            .withExistingPrimaryNetwork(network)
+            .withSubnet(generateRandomResourceName("subnet", 15))
+            .withPrimaryPrivateIPAddressDynamic()
+            .withoutPrimaryPublicIPAddress()
+            .withPopularWindowsImage(KnownWindowsVirtualMachineImage.WINDOWS_SERVER_2019_DATACENTER_GEN2)
+            .withAdminUsername("Foo12")
+            .withAdminPassword(password())
+            .withNewDataDisk(127)
+            .withSize(VirtualMachineSizeTypes.STANDARD_DS1_V2)
+            .withCapacityReservationGroup(crgForCreate.id())
+            .create();
+
+        Assertions.assertEquals(crgForCreate.id(), vm.capacityReservationGroupId());
+
+        // Update virtual machine with another capacity reservation group
+        vm.update().withCapacityReservationGroup(crgForUpdate.id()).apply();
+
+        Assertions.assertEquals(crgForUpdate.id(), vm.capacityReservationGroupId());
+    }
+
     @Test
     public void canCreateAndUpdateVirtualMachineWithUserData() {
         String userDataForCreate = "N0ZBN0MxRkYtMkNCMC00RUM3LUE1RDctMDY2MUI0RTdDNzY4";