Merge branch 'main' into tpu-vm-crud-operations

TetyanaYahodska · TetyanaYahodska · commit b804cc85e47e · 2024-10-29T14:10:44.000+01:00
diff --git a/compute/cloud-client/pom.xml b/compute/cloud-client/pom.xml
@@ -59,6 +59,13 @@
       <groupId>com.google.cloud</groupId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <version>5.13.0</version>
+      <scope>test</scope>
+    </dependency>
+
 
     <dependency>
       <artifactId>truth</artifactId>
diff --git a/compute/cloud-client/src/main/java/compute/reservation/CreateSharedReservation.java b/compute/cloud-client/src/main/java/compute/reservation/CreateSharedReservation.java
@@ -0,0 +1,110 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package compute.reservation;
+
+// [START compute_reservation_create_shared]
+
+import com.google.cloud.compute.v1.AllocationSpecificSKUReservation;
+import com.google.cloud.compute.v1.Operation;
+import com.google.cloud.compute.v1.Reservation;
+import com.google.cloud.compute.v1.ReservationsClient;
+import com.google.cloud.compute.v1.ShareSettings;
+import com.google.cloud.compute.v1.ShareSettingsProjectConfig;
+import java.io.IOException;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+public class CreateSharedReservation {
+  private final ReservationsClient reservationsClient;
+
+  // Constructor to inject the ReservationsClient
+  public CreateSharedReservation(ReservationsClient reservationsClient) {
+    this.reservationsClient = reservationsClient;
+  }
+
+  public static void main(String[] args)
+      throws IOException, ExecutionException, InterruptedException, TimeoutException {
+    // TODO(developer): Replace these variables before running the sample.
+    // The ID of the project where you want to reserve resources
+    // and where the instance template exists.
+    // By default, no projects are allowed to create or modify shared reservations
+    // in an organization. Add projects to the Shared Reservations Owner Projects
+    // (compute.sharedReservationsOwnerProjects) organization policy constraint
+    // to allow them to create and modify shared reservations.
+    // For more information visit this page:
+    // https://cloud.google.com/compute/docs/instances/reservations-shared#shared_reservation_constraint
+    String projectId = "YOUR_PROJECT_ID";
+    // Zone in which the reservation resides.
+    String zone = "us-central1-a";
+    // Name of the reservation to be created.
+    String reservationName = "YOUR_RESERVATION_NAME";
+    // The URI of the global instance template to be used for creating the reservation.
+    String instanceTemplateUri = String.format(
+        "projects/%s/global/instanceTemplates/YOUR_INSTANCE_TEMPLATE_NAME", projectId);
+    // Number of instances for which capacity needs to be reserved.
+    int vmCount = 3;
+    // In your main method, create ReservationsClient
+    ReservationsClient client = ReservationsClient.create();
+    // Create an instance of your class, passing in the client
+    CreateSharedReservation creator = new CreateSharedReservation(client);
+
+    creator.createSharedReservation(projectId, zone, reservationName, instanceTemplateUri, vmCount);
+  }
+
+  // Creates a shared reservation with the given name in the given zone.
+  public void createSharedReservation(
+      String projectId, String zone,
+      String reservationName, String instanceTemplateUri, int vmCount)
+      throws ExecutionException, InterruptedException, TimeoutException {
+
+    ShareSettings shareSettings = ShareSettings.newBuilder()
+        .setShareType(String.valueOf(ShareSettings.ShareType.SPECIFIC_PROJECTS))
+        // The IDs of projects that can consume this reservation. You can include up to 100
+        // consumer projects. These projects must be in the same organization as
+        // the owner project. Don't include the owner project. By default, it is already allowed
+        // to consume the reservation.
+        .putProjectMap("CONSUMER_PROJECT_ID_1", ShareSettingsProjectConfig.newBuilder().build())
+        .putProjectMap("CONSUMER_PROJECT_ID_2", ShareSettingsProjectConfig.newBuilder().build())
+        .build();
+
+    // Create the reservation.
+    Reservation reservation =
+        Reservation.newBuilder()
+            .setName(reservationName)
+            .setZone(zone)
+            .setSpecificReservationRequired(true)
+            .setShareSettings(shareSettings)
+            .setSpecificReservation(
+                AllocationSpecificSKUReservation.newBuilder()
+                    .setCount(vmCount)
+                    .setSourceInstanceTemplate(instanceTemplateUri)
+                    .build())
+            .build();
+
+    // Wait for the create reservation operation to complete.
+    Operation response =
+        this.reservationsClient.insertAsync(projectId, zone, reservation).get(3, TimeUnit.MINUTES);
+
+    if (response.hasError()) {
+      System.out.println("Reservation creation failed!" + response);
+      return;
+    }
+    System.out.println("Reservation created. Operation Status: " + response.getStatus());
+  }
+}
+// [END compute_reservation_create_shared]
diff --git a/compute/cloud-client/src/test/java/compute/reservation/ReservationIT.java b/compute/cloud-client/src/test/java/compute/reservation/ReservationIT.java
@@ -18,10 +18,19 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+import com.google.api.gax.longrunning.OperationFuture;
 import com.google.api.gax.rpc.NotFoundException;
+import com.google.cloud.compute.v1.AllocationSpecificSKUReservation;
+import com.google.cloud.compute.v1.Operation;
 import com.google.cloud.compute.v1.Reservation;
 import com.google.cloud.compute.v1.ReservationsClient;
+import com.google.cloud.compute.v1.ShareSettings;
+import com.google.cloud.compute.v1.ShareSettingsProjectConfig;
 import compute.CreateInstanceTemplate;
 import compute.CreateRegionalInstanceTemplate;
 import compute.DeleteInstanceTemplate;
@@ -36,8 +45,10 @@
 import java.util.concurrent.TimeoutException;
 import org.junit.Assert;
 import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.MethodOrderer;
 import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
@@ -69,7 +80,16 @@ public class ReservationIT {
   private static final String REGIONAL_INSTANCE_TEMPLATE_URI =
       String.format("projects/%s/regions/%s/instanceTemplates/%s",
           PROJECT_ID, REGION, REGIONAL_INSTANCE_TEMPLATE_NAME);
+  private static final String SPECIFIC_SHARED_INSTANCE_TEMPLATE_NAME =
+      "test-shared-inst-temp-"  + javaVersion  + "-"
+          + UUID.randomUUID().toString().substring(0, 8);
+  private static final String  INSTANCE_TEMPLATE_SHARED_RESERV_URI =
+      String.format("projects/%s/global/instanceTemplates/%s",
+      PROJECT_ID, SPECIFIC_SHARED_INSTANCE_TEMPLATE_NAME);
+  private static final String RESERVATION_NAME_SHARED = "test-reservation-shared-" + javaVersion
+      + "-" + UUID.randomUUID().toString().substring(0, 8);
   private static final int NUMBER_OF_VMS = 3;
+  private ByteArrayOutputStream stdOut;
 
   // Check if the required environment variables are set.
   public static void requireEnvVar(String envVarName) {
@@ -93,6 +113,7 @@ public static void setUp()
     Util.cleanUpExistingReservations(
         "test-reservation-global-" + javaVersion, PROJECT_ID, ZONE);
     Util.cleanUpExistingReservations("test-reservation-regional-" + javaVersion, PROJECT_ID, ZONE);
+    Util.cleanUpExistingInstanceTemplates("test-shared-inst-temp-" + javaVersion, PROJECT_ID);
 
     // Initialize the client once for all tests
     reservationsClient = ReservationsClient.create();
@@ -105,6 +126,9 @@ public static void setUp()
     CreateRegionalInstanceTemplate.createRegionalInstanceTemplate(
         PROJECT_ID, REGION, REGIONAL_INSTANCE_TEMPLATE_NAME);
     assertThat(stdOut.toString()).contains("Instance Template Operation Status: DONE");
+    // Create instance template for shares reservation.
+    CreateInstanceTemplate.createInstanceTemplate(
+        PROJECT_ID, SPECIFIC_SHARED_INSTANCE_TEMPLATE_NAME);
 
     stdOut.close();
     System.setOut(out);
@@ -130,6 +154,13 @@ public static void cleanup()
         .contains("Instance template deletion operation status for "
             + REGIONAL_INSTANCE_TEMPLATE_NAME);
 
+    // Delete instance template for shared reservation
+    DeleteInstanceTemplate.deleteInstanceTemplate(
+        PROJECT_ID, SPECIFIC_SHARED_INSTANCE_TEMPLATE_NAME);
+    assertThat(stdOut.toString())
+        .contains("Instance template deletion operation status for "
+            + SPECIFIC_SHARED_INSTANCE_TEMPLATE_NAME);
+
     // Delete all reservations created for testing.
     DeleteReservation.deleteReservation(PROJECT_ID, ZONE, RESERVATION_NAME_GLOBAL);
     DeleteReservation.deleteReservation(PROJECT_ID, ZONE, RESERVATION_NAME_REGIONAL);
@@ -149,6 +180,18 @@ public static void cleanup()
     System.setOut(out);
   }
 
+  @BeforeEach
+  public void beforeEach() {
+    stdOut = new ByteArrayOutputStream();
+    System.setOut(new PrintStream(stdOut));
+  }
+
+  @AfterEach
+  public void afterEach() {
+    stdOut = null;
+    System.setOut(null);
+  }
+
   @Test
   @Order(1)
   public void testCreateReservationWithGlobalInstanceTemplate()
@@ -189,4 +232,53 @@ public void testUpdateVmsForReservation()
 
     Assert.assertEquals(newNumberOfVms, reservation.getSpecificReservation().getCount());
   }
+
+  @Test
+  public void testCreateSharedReservation()
+      throws ExecutionException, InterruptedException, TimeoutException {
+    // Mock the ReservationsClient
+    ReservationsClient mockReservationsClient = mock(ReservationsClient.class);
+
+    // This test require projects in the test environment to share reservation with,
+    // therefore the operation should be mocked. If you want to make a real test,
+    // please set the CONSUMER_PROJECT_ID_1 and CONSUMER_PROJECT_ID_2 accordingly.
+    // Make sure that base project has proper permissions to share reservations.
+    // See: https://cloud.google.com/compute/docs/instances/reservations-shared#shared_reservation_constraint
+    ShareSettings shareSettings = ShareSettings.newBuilder()
+        .setShareType(String.valueOf(ShareSettings.ShareType.SPECIFIC_PROJECTS))
+        .putProjectMap("CONSUMER_PROJECT_ID_1", ShareSettingsProjectConfig.newBuilder().build())
+        .putProjectMap("CONSUMER_PROJECT_ID_2", ShareSettingsProjectConfig.newBuilder().build())
+        .build();
+
+    Reservation reservation =
+        Reservation.newBuilder()
+            .setName(RESERVATION_NAME_SHARED)
+            .setZone(ZONE)
+            .setSpecificReservationRequired(true)
+            .setShareSettings(shareSettings)
+            .setSpecificReservation(
+                AllocationSpecificSKUReservation.newBuilder()
+                    .setCount(NUMBER_OF_VMS)
+                    .setSourceInstanceTemplate(INSTANCE_TEMPLATE_SHARED_RESERV_URI)
+                    .build())
+            .build();
+
+    OperationFuture mockFuture = mock(OperationFuture.class);
+    when(mockReservationsClient.insertAsync(PROJECT_ID, ZONE, reservation))
+        .thenReturn(mockFuture);
+    Operation mockOperation = mock(Operation.class);
+    when(mockFuture.get(3, TimeUnit.MINUTES)).thenReturn(mockOperation);
+    when(mockOperation.hasError()).thenReturn(false);
+    when(mockOperation.getStatus()).thenReturn(Operation.Status.DONE);
+
+    // Create an instance, passing in the mock client
+    CreateSharedReservation creator = new CreateSharedReservation(mockReservationsClient);
+
+    creator.createSharedReservation(PROJECT_ID, ZONE,
+        RESERVATION_NAME_SHARED, INSTANCE_TEMPLATE_SHARED_RESERV_URI, NUMBER_OF_VMS);
+
+    verify(mockReservationsClient, times(1))
+        .insertAsync(PROJECT_ID, ZONE, reservation);
+    assertThat(stdOut.toString()).contains("Reservation created. Operation Status: DONE");
+  }
 }
