Added comment, created Util class

Author: TetyanaYahodska

tpu/src/main/java/tpu/DeleteTpuVm.java

@@ -48,6 +48,8 @@ public static void main(String[] args)
   // Deletes a TPU VM with the specified name in the given project and zone.
   public static void deleteTpuVm(String projectId, String zone, String nodeName)
       throws IOException, ExecutionException, InterruptedException {
+    // With these settings the client library handles the Operation's polling mechanism
+    // and prevent CancellationException error
     TpuSettings.Builder clientSettings =
         TpuSettings.newBuilder();
     clientSettings

tpu/src/test/java/tpu/TpuVmIT.java

@@ -25,15 +25,9 @@
 import com.google.api.gax.rpc.NotFoundException;
 import com.google.cloud.tpu.v2.Node;
 import com.google.cloud.tpu.v2.TpuClient;
-import com.google.protobuf.Timestamp;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintStream;
-import java.time.Instant;
-import java.time.OffsetDateTime;
-import java.time.ZoneOffset;
-import java.time.format.DateTimeFormatter;
-import java.time.temporal.ChronoUnit;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
@@ -75,13 +69,12 @@ public static void setUp()
     requireEnvVar("GOOGLE_CLOUD_PROJECT");
 
     // Cleanup existing stale resources.
-    cleanUpExistingTpu("test-tpu-" + javaVersion, PROJECT_ID, ZONE);
+    Util.cleanUpExistingTpu("test-tpu-" + javaVersion, PROJECT_ID, ZONE);
   }
 
   @AfterAll
   public static void cleanup() throws Exception {
     DeleteTpuVm.deleteTpuVm(PROJECT_ID, ZONE, TPU_VM_NAME);
-    TimeUnit.MINUTES.sleep(5);
 
     // Test that TPUs is deleted
     Assertions.assertThrows(
@@ -96,7 +89,6 @@ public void testCreateTpuVm() throws IOException, ExecutionException, Interrupte
     ByteArrayOutputStream stdOut = new ByteArrayOutputStream();
     System.setOut(new PrintStream(stdOut));
     CreateTpuVm.createTpuVm(PROJECT_ID, ZONE, TPU_VM_NAME, ACCELERATOR_TYPE, VERSION);
-    TimeUnit.MINUTES.sleep(3);
 
     assertThat(stdOut.toString()).contains("TPU VM created: " + TPU_VM_PATH_NAME);
     stdOut.close();
@@ -140,45 +132,4 @@ public void testStartTpuVm() throws IOException, ExecutionException, Interrupted
 
     assertThat(node.getState()).isEqualTo(READY);
   }
-
-  public static void cleanUpExistingTpu(String prefixToDelete, String projectId, String zone)
-      throws IOException, ExecutionException, InterruptedException {
-    try (TpuClient tpuClient = TpuClient.create()) {
-      String parent = String.format("projects/%s/locations/%s", projectId, zone);
-      for (Node node : tpuClient.listNodes(parent).iterateAll()) {
-        String creationTime = formatTimestamp(node.getCreateTime());
-        String name = node.getName().substring(node.getName().lastIndexOf("/") + 1);
-        if (containPrefixToDeleteAndZone(node, prefixToDelete, zone)
-            && isCreatedBeforeThresholdTime(creationTime)) {
-          DeleteTpuVm.deleteTpuVm(projectId, zone, name);
-        }
-      }
-    }
-  }
-
-  public static boolean containPrefixToDeleteAndZone(
-      Node node, String prefixToDelete, String zone) {
-    boolean containPrefixAndZone = false;
-    try {
-      containPrefixAndZone = node.getName().contains(prefixToDelete)
-          && node.getName().split("/")[3].contains(zone);
-
-    } catch (NullPointerException e) {
-      System.out.println("Resource not found, skipping deletion:");
-    }
-    return containPrefixAndZone;
-  }
-
-  public static boolean isCreatedBeforeThresholdTime(String timestamp) {
-    return OffsetDateTime.parse(timestamp).toInstant()
-        .isBefore(Instant.now().minus(30, ChronoUnit.MINUTES));
-  }
-
-  private static String formatTimestamp(Timestamp timestamp) {
-    DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSSXXX");
-    OffsetDateTime offsetDateTime = OffsetDateTime.ofInstant(
-        Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos()),
-        ZoneOffset.UTC);
-    return formatter.format(offsetDateTime);
-  }
 }

tpu/src/test/java/tpu/Util.java

@@ -0,0 +1,75 @@
+/*
+ * 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 tpu;
+
+import com.google.cloud.tpu.v2.Node;
+import com.google.cloud.tpu.v2.TpuClient;
+import com.google.protobuf.Timestamp;
+import java.io.IOException;
+import java.time.Instant;
+import java.time.OffsetDateTime;
+import java.time.ZoneOffset;
+import java.time.format.DateTimeFormatter;
+import java.time.temporal.ChronoUnit;
+import java.util.concurrent.ExecutionException;
+
+public class Util {
+  private static final int DELETION_THRESHOLD_TIME_MINUTES = 30;
+
+  // Delete TPU VMs which starts with the given prefixToDelete and
+  // has creation timestamp >30 minutes.
+  public static void cleanUpExistingTpu(String prefixToDelete, String projectId, String zone)
+      throws IOException, ExecutionException, InterruptedException {
+    try (TpuClient tpuClient = TpuClient.create()) {
+      String parent = String.format("projects/%s/locations/%s", projectId, zone);
+      for (Node node : tpuClient.listNodes(parent).iterateAll()) {
+        String creationTime = formatTimestamp(node.getCreateTime());
+        String name = node.getName().substring(node.getName().lastIndexOf("/") + 1);
+        if (containPrefixToDeleteAndZone(node, prefixToDelete, zone)
+            && isCreatedBeforeThresholdTime(creationTime)) {
+          DeleteTpuVm.deleteTpuVm(projectId, zone, name);
+        }
+      }
+    }
+  }
+
+  public static boolean containPrefixToDeleteAndZone(
+      Node node, String prefixToDelete, String zone) {
+    boolean containPrefixAndZone = false;
+    try {
+      containPrefixAndZone = node.getName().contains(prefixToDelete)
+          && node.getName().split("/")[3].contains(zone);
+
+    } catch (NullPointerException e) {
+      System.out.println("Resource not found, skipping deletion:");
+    }
+    return containPrefixAndZone;
+  }
+
+  public static boolean isCreatedBeforeThresholdTime(String timestamp) {
+    return OffsetDateTime.parse(timestamp).toInstant()
+        .isBefore(Instant.now().minus(DELETION_THRESHOLD_TIME_MINUTES, ChronoUnit.MINUTES));
+  }
+
+  private static String formatTimestamp(Timestamp timestamp) {
+    DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSSXXX");
+    OffsetDateTime offsetDateTime = OffsetDateTime.ofInstant(
+        Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos()),
+        ZoneOffset.UTC);
+    return formatter.format(offsetDateTime);
+  }
+}