Fixed tests

Author: TetyanaYahodska

tpu/src/main/java/tpu/CreateTimeBoundQueuedResource.java

@@ -24,12 +24,12 @@
 import com.google.cloud.tpu.v2alpha1.TpuClient;
 import com.google.cloud.tpu.v2alpha1.TpuSettings;
 import com.google.protobuf.Duration;
-import com.google.protobuf.Timestamp;
 // Uncomment the following line to use Interval or Date
+//import com.google.protobuf.Timestamp;
 //import com.google.type.Interval;
 //import java.util.Date;
+//import java.time.Instant;
 import java.io.IOException;
-import java.time.Instant;
 import java.util.concurrent.ExecutionException;
 
 public class CreateTimeBoundQueuedResource {
@@ -161,21 +161,20 @@ public static QueuedResource createTimeBoundQueuedResource(
               .setQueuedResourceId(queuedResourceName)
               .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", response.getName());
-      return response;
+      return tpuClient.createQueuedResourceAsync(request).get();
     }
   }
 
+  // Uncomment this method if you want to use time bound configurations
   // Method to convert a string timestamp to a Protobuf Timestamp object
-  private static Timestamp convertStringToTimestamp(String timestampString) {
-    Instant instant = Instant.parse(timestampString);
-    return Timestamp.newBuilder()
-        .setSeconds(instant.getEpochSecond())
-        .setNanos(instant.getNano())
-        .build();
-  }
+  //  private static Timestamp convertStringToTimestamp(String timestampString) {
+  //    Instant instant = Instant.parse(timestampString);
+  //    return Timestamp.newBuilder()
+  //        .setSeconds(instant.getEpochSecond())
+  //        .setNanos(instant.getNano())
+  //        .build();
+  //  }
 }
 // [END tpu_queued_resources_time_bound]

tpu/src/test/java/tpu/CreateTimeBoundQueuedResourceIT.java

@@ -35,15 +35,15 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintStream;
-import org.junit.Before;
-import org.junit.Test;
+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;
 import org.mockito.MockedStatic;
 
 @RunWith(JUnit4.class)
-@Timeout(value = 3)
+@Timeout(value = 30)
 public class QueuedResourceIT {
   private static final String PROJECT_ID = "project-id";
   private static final String ZONE = "europe-west4-a";
@@ -52,10 +52,10 @@ public class QueuedResourceIT {
   private static final String TPU_SOFTWARE_VERSION = "tpu-vm-tf-2.14.1";
   private static final String QUEUED_RESOURCE_NAME = "queued-resource";
   private static final String NETWORK_NAME = "default";
-  private ByteArrayOutputStream bout;
+  private static ByteArrayOutputStream bout;
 
-  @Before
-  public void setUp() {
+  @BeforeAll
+  public static void setUp() {
     bout = new ByteArrayOutputStream();
     System.setOut(new PrintStream(bout));
   }
@@ -75,8 +75,8 @@ public void testCreateQueuedResourceWithSpecifiedNetwork() throws Exception {
 
       QueuedResource returnedQueuedResource =
           CreateQueuedResourceWithNetwork.createQueuedResourceWithNetwork(
-          PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME, NODE_NAME,
-          TPU_TYPE, TPU_SOFTWARE_VERSION, NETWORK_NAME);
+              PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME, NODE_NAME,
+              TPU_TYPE, TPU_SOFTWARE_VERSION, NETWORK_NAME);
 
       verify(mockTpuClient, times(1))
           .createQueuedResourceAsync(any(CreateQueuedResourceRequest.class));
@@ -89,7 +89,6 @@ public void testCreateQueuedResourceWithSpecifiedNetwork() throws Exception {
   public void testGetQueuedResource() throws IOException {
     try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
       TpuClient mockClient = mock(TpuClient.class);
-      GetQueuedResource mockGetQueuedResource = mock(GetQueuedResource.class);
       QueuedResource mockQueuedResource = mock(QueuedResource.class);
 
       mockedTpuClient.when(TpuClient::create).thenReturn(mockClient);
@@ -99,14 +98,14 @@ public void testGetQueuedResource() throws IOException {
       QueuedResource returnedQueuedResource =
           GetQueuedResource.getQueuedResource(PROJECT_ID, ZONE, NODE_NAME);
 
-      verify(mockGetQueuedResource, times(1))
-          .getQueuedResource(PROJECT_ID, ZONE, NODE_NAME);
+      verify(mockClient, times(1))
+          .getQueuedResource(any(GetQueuedResourceRequest.class));
       assertEquals(returnedQueuedResource, mockQueuedResource);
     }
   }
 
   @Test
-  public void testDeleteTpuVm() {
+  public void testDeleteForceQueuedResource() {
     try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
       TpuClient mockTpuClient = mock(TpuClient.class);
       OperationFuture mockFuture = mock(OperationFuture.class);
@@ -124,4 +123,29 @@ public void testDeleteTpuVm() {
           .deleteQueuedResourceAsync(any(DeleteQueuedResourceRequest.class));
     }
   }
+
+  @Test
+  public void testCreateTimeBoundQueuedResource() throws Exception {
+    try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
+      QueuedResource mockQueuedResource = mock(QueuedResource.class);
+      TpuClient mockTpuClient = mock(TpuClient.class);
+      OperationFuture mockFuture = mock(OperationFuture.class);
+
+      mockedTpuClient.when(() -> TpuClient.create(any(TpuSettings.class)))
+          .thenReturn(mockTpuClient);
+      when(mockTpuClient.createQueuedResourceAsync(any(CreateQueuedResourceRequest.class)))
+          .thenReturn(mockFuture);
+      when(mockFuture.get()).thenReturn(mockQueuedResource);
+
+      QueuedResource returnedQueuedResource =
+          CreateTimeBoundQueuedResource.createTimeBoundQueuedResource(
+              PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME, NODE_NAME,
+              TPU_TYPE, TPU_SOFTWARE_VERSION);
+
+      verify(mockTpuClient, times(1))
+          .createQueuedResourceAsync(any(CreateQueuedResourceRequest.class));
+      verify(mockFuture, times(1)).get();
+      assertEquals(returnedQueuedResource, mockQueuedResource);
+    }
+  }
 }