Fixed test, deleted cleanup method

Author: TetyanaYahodska

tpu/src/main/java/tpu/CreateQueuedResourceWithNetwork.java

@@ -127,12 +127,10 @@ public static QueuedResource createQueuedResourceWithNetwork(
               .setQueuedResource(queuedResource)
               .setQueuedResourceId(queuedResourceId)
               .build();
-
-      QueuedResource response = tpuClient.createQueuedResourceAsync(request).get();
       // You can wait until TPU Node is READY,
       // and check its status using getTpuVm() from "tpu_vm_get" sample.
-      System.out.println("Queued Resource created: " + queuedResourceId);
-      return response;
+
+      return tpuClient.createQueuedResourceAsync(request).get();
     }
   }
 }

tpu/src/main/java/tpu/CreateQueuedResourceWithStartupScript.java

@@ -24,8 +24,6 @@
 import com.google.cloud.tpu.v2alpha1.TpuClient;
 import com.google.cloud.tpu.v2alpha1.TpuSettings;
 import java.io.IOException;
-import java.nio.file.Files;
-import java.nio.file.Paths;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
@@ -52,18 +50,15 @@ public static void main(String[] args)
     String tpuSoftwareVersion = "tpu-vm-tf-2.14.1";
     // The name for your Queued Resource.
     String queuedResourceId = "QUEUED_RESOURCE_ID";
-    // Path to the startup script file.
-    String startupScriptPath = "YOUR_STARTUP_SCRIPT_PATH";
 
-    createQueuedResource(
-        projectId, zone, queuedResourceId, nodeName,
-        tpuType, tpuSoftwareVersion, startupScriptPath);
+    createQueuedResource(projectId, zone, queuedResourceId, nodeName,
+        tpuType, tpuSoftwareVersion);
   }
 
   // Creates a Queued Resource with startup script.
   public static QueuedResource createQueuedResource(
       String projectId, String zone, String queuedResourceId,
-      String nodeName, String tpuType, String tpuSoftwareVersion, String startupScriptPath)
+      String nodeName, String tpuType, String tpuSoftwareVersion)
       throws IOException, ExecutionException, InterruptedException {
     // With these settings the client library handles the Operation's polling mechanism
     // and prevent CancellationException error
@@ -80,13 +75,14 @@ public static QueuedResource createQueuedResource(
                 .setMaxRetryDelay(Duration.ofMillis(45000L))
                 .setTotalTimeout(Duration.ofHours(24L))
                 .build());
-      String parent = String.format("projects/%s/locations/%s", projectId, zone);
-      // Read the startup script content from the file
-      String startupScriptContent = new String(Files.readAllBytes(Paths.get(startupScriptPath)));
-      // Add startup script to metadata
-      Map<String, String> metadata = new HashMap<>();
-      metadata.put("startup-script", startupScriptContent);
-
+    String parent = String.format("projects/%s/locations/%s", projectId, zone);
+    // Read the startup script content from the file
+    String startupScriptContent = "#!/bin/bash\necho \"Hello from the startup script!\"";
+    // Add startup script to metadata
+    Map<String, String> metadata = new HashMap<>();
+    metadata.put("startup-script", startupScriptContent);
+    String queuedResourceForTpu =  String.format("projects/%s/locations/%s/queuedResources/%s",
+            projectId, zone, queuedResourceId);
     // Initialize client that will be used to send requests. This client only needs to be created
     // once, and can be reused for multiple requests.
     try (TpuClient tpuClient = TpuClient.create(clientSettings.build())) {
@@ -95,10 +91,7 @@ public static QueuedResource createQueuedResource(
               .setName(nodeName)
               .setAcceleratorType(tpuType)
               .setRuntimeVersion(tpuSoftwareVersion)
-              .setQueuedResource(
-                  String.format(
-                      "projects/%s/locations/%s/queuedResources/%s",
-                      projectId, zone, queuedResourceId))
+              .setQueuedResource(queuedResourceForTpu)
               .putAllMetadata(metadata)
               .build();
 
@@ -122,12 +115,10 @@ public static QueuedResource createQueuedResource(
               .setQueuedResourceId(queuedResourceId)
               .setQueuedResource(queuedResource)
               .build();
-
-      QueuedResource response = tpuClient.createQueuedResourceAsync(request).get();
       // You can wait until TPU Node is READY,
       // and check its status using getTpuVm() from "tpu_vm_get" sample.
-      System.out.printf("Queued Resource created: %s\n", queuedResourceId);
-      return response;
+
+      return tpuClient.createQueuedResourceAsync(request).get();
     }
   }
 }

tpu/src/main/java/tpu/DeleteForceQueuedResource.java

@@ -17,7 +17,6 @@
 package tpu;
 
 //[START tpu_queued_resources_delete_force]
-
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.rpc.UnknownException;
 import com.google.cloud.tpu.v2alpha1.DeleteQueuedResourceRequest;

tpu/src/test/java/tpu/CreateQueuedResourceWithNetworkIT.java

@@ -19,29 +19,33 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
-import com.google.api.gax.rpc.NotFoundException;
 import com.google.cloud.tpu.v2alpha1.QueuedResource;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
+import org.junit.Test;
 import org.junit.jupiter.api.AfterAll;
-import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 @RunWith(JUnit4.class)
 @Timeout(value = 6, unit = TimeUnit.MINUTES)
-public class CreateQueuedResourceWithStartupScriptIT {
+public class QueuedResourceIT {
 
   private static final String PROJECT_ID = System.getenv("GOOGLE_CLOUD_PROJECT");
   private static final String ZONE = "europe-west4-a";
-  private static final String NODE_NAME = "test-tpu-queued-resource-script-" + UUID.randomUUID();
+  private static final String NODE_FOR_NETWORK =
+      "test-tpu-queued-resource-network-" + UUID.randomUUID();
+  private static final String NODE_FOR_SCRIPT =
+      "test-tpu-queued-resource-script-" + UUID.randomUUID();
   private static final String TPU_TYPE = "v2-8";
   private static final String TPU_SOFTWARE_VERSION = "tpu-vm-tf-2.14.1";
-  private static final String QUEUED_RESOURCE_NAME = "queued-resource-script-" + UUID.randomUUID();
-  private static final String STARTUP_SCRIPT_PATH = "src/test/java/tpu/startup-script.sh";
+  private static final String QUEUED_RESOURCE_FOR_NETWORK =
+      "queued-resource-network-" + UUID.randomUUID();
+  private static final String QUEUED_RESOURCE_FOR_SCRIPT =
+      "queued-resource-script-" + UUID.randomUUID();
+  private static final String NETWORK_NAME = "default";
 
   public static void requireEnvVar(String envVarName) {
     assertWithMessage(String.format("Missing environment variable '%s' ", envVarName))
@@ -56,21 +60,35 @@ public static void setUp() {
 
   @AfterAll
   public static void cleanup() {
-    DeleteForceQueuedResource.deleteForceQueuedResource(PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME);
+    DeleteForceQueuedResource.deleteForceQueuedResource(
+        PROJECT_ID, ZONE, QUEUED_RESOURCE_FOR_NETWORK);
+    DeleteForceQueuedResource.deleteForceQueuedResource(
+        PROJECT_ID, ZONE, QUEUED_RESOURCE_FOR_SCRIPT);
+  }
+
+  @Test
+  public void testCreateQueuedResourceWithSpecifiedNetwork() throws Exception {
+
+    QueuedResource queuedResource = CreateQueuedResourceWithNetwork.createQueuedResourceWithNetwork(
+        PROJECT_ID, ZONE, QUEUED_RESOURCE_FOR_NETWORK, NODE_FOR_NETWORK,
+        TPU_TYPE, TPU_SOFTWARE_VERSION, NETWORK_NAME);
 
-    Assertions.assertThrows(
-        NotFoundException.class,
-        () -> GetQueuedResource.getQueuedResource(PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME));
+    assertThat(queuedResource.getTpu().getNodeSpec(0).getNode().getName())
+        .isEqualTo(NODE_FOR_NETWORK);
+    assertThat(queuedResource.getTpu().getNodeSpec(0).getNode().getNetworkConfig().getNetwork()
+        .contains(NETWORK_NAME));
+    assertThat(queuedResource.getTpu().getNodeSpec(0).getNode().getNetworkConfig().getSubnetwork()
+        .contains(NETWORK_NAME));
   }
 
   @Test
   public void testCreateQueuedResourceWithStartupScript() throws Exception {
     QueuedResource queuedResource = CreateQueuedResourceWithStartupScript.createQueuedResource(
-        PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME, NODE_NAME,
-        TPU_TYPE, TPU_SOFTWARE_VERSION, STARTUP_SCRIPT_PATH);
+        PROJECT_ID, ZONE, QUEUED_RESOURCE_FOR_SCRIPT, NODE_FOR_SCRIPT,
+        TPU_TYPE, TPU_SOFTWARE_VERSION);
 
     assertThat(queuedResource.getTpu().getNodeSpec(0).getNode().containsLabels("startup-script"));
     assertThat(queuedResource.getTpu().getNodeSpec(0).getNode().getLabelsMap()
-        .containsValue("your-script-here"));
+        .containsValue("Hello from the startup script!"));
   }
-}
+}