Fixed tests

Author: TetyanaYahodska

tpu/src/test/java/tpu/CreateTpuIT.java

@@ -16,6 +16,7 @@
 
 package tpu;
 
+import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.mockStatic;
@@ -25,6 +26,7 @@
 
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.cloud.tpu.v2.CreateNodeRequest;
+import com.google.cloud.tpu.v2.Node;
 import com.google.cloud.tpu.v2.TpuClient;
 import com.google.cloud.tpu.v2.TpuSettings;
 import org.junit.Test;
@@ -45,19 +47,24 @@ public class CreateTpuIT {
   @Test
   public void testCreateTpuVm() throws Exception {
     try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
+      Node mockNode = mock(Node.class);
       TpuClient mockTpuClient = mock(TpuClient.class);
+      OperationFuture mockFuture = mock(OperationFuture.class);
+
       mockedTpuClient.when(() -> TpuClient.create(any(TpuSettings.class)))
           .thenReturn(mockTpuClient);
-
-      OperationFuture mockFuture = mock(OperationFuture.class);
       when(mockTpuClient.createNodeAsync(any(CreateNodeRequest.class)))
           .thenReturn(mockFuture);
-      CreateTpuVm.createTpuVm(
+      when(mockFuture.get()).thenReturn(mockNode);
+
+      Node returnedNode = CreateTpuVm.createTpuVm(
           PROJECT_ID, ZONE, NODE_NAME,
           TPU_TYPE, TPU_SOFTWARE_VERSION);
 
       verify(mockTpuClient, times(1))
           .createNodeAsync(any(CreateNodeRequest.class));
+      verify(mockFuture, times(1)).get();
+      assertEquals(returnedNode, mockNode);
     }
   }
 }

tpu/src/test/java/tpu/QueuedResourceIT.java

@@ -17,6 +17,7 @@
 package tpu;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.mockStatic;
@@ -27,8 +28,8 @@
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.cloud.tpu.v2alpha1.CreateQueuedResourceRequest;
 import com.google.cloud.tpu.v2alpha1.DeleteQueuedResourceRequest;
+import com.google.cloud.tpu.v2alpha1.GetQueuedResourceRequest;
 import com.google.cloud.tpu.v2alpha1.QueuedResource;
-import com.google.cloud.tpu.v2alpha1.QueuedResourceName;
 import com.google.cloud.tpu.v2alpha1.TpuClient;
 import com.google.cloud.tpu.v2alpha1.TpuSettings;
 import java.io.ByteArrayOutputStream;
@@ -62,48 +63,59 @@ public void setUp() {
   @Test
   public void testCreateQueuedResourceWithSpecifiedNetwork() 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);
-
-      OperationFuture mockFuture = mock(OperationFuture.class);
       when(mockTpuClient.createQueuedResourceAsync(any(CreateQueuedResourceRequest.class)))
           .thenReturn(mockFuture);
-      CreateQueuedResourceWithNetwork.createQueuedResourceWithNetwork(
+      when(mockFuture.get()).thenReturn(mockQueuedResource);
+
+      QueuedResource returnedQueuedResource =
+          CreateQueuedResourceWithNetwork.createQueuedResourceWithNetwork(
           PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME, NODE_NAME,
           TPU_TYPE, TPU_SOFTWARE_VERSION, NETWORK_NAME);
 
       verify(mockTpuClient, times(1))
           .createQueuedResourceAsync(any(CreateQueuedResourceRequest.class));
+      verify(mockFuture, times(1)).get();
+      assertEquals(returnedQueuedResource, mockQueuedResource);
     }
   }
 
   @Test
   public void testGetQueuedResource() throws IOException {
     try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
-      QueuedResource mockQueuedResource = mock(QueuedResource.class);
-      mockedTpuClient.when(TpuClient::create).thenReturn(mock(TpuClient.class));
-      when(mock(TpuClient.class)
-          .getQueuedResource(any(QueuedResourceName.class))).thenReturn(mockQueuedResource);
+      TpuClient mockClient = mock(TpuClient.class);
       GetQueuedResource mockGetQueuedResource = mock(GetQueuedResource.class);
+      QueuedResource mockQueuedResource = mock(QueuedResource.class);
+
+      mockedTpuClient.when(TpuClient::create).thenReturn(mockClient);
+      when(mockClient.getQueuedResource(any(GetQueuedResourceRequest.class)))
+          .thenReturn(mockQueuedResource);
 
-      GetQueuedResource.getQueuedResource(PROJECT_ID, ZONE, NODE_NAME);
+      QueuedResource returnedQueuedResource =
+          GetQueuedResource.getQueuedResource(PROJECT_ID, ZONE, NODE_NAME);
 
       verify(mockGetQueuedResource, times(1))
           .getQueuedResource(PROJECT_ID, ZONE, NODE_NAME);
+      assertEquals(returnedQueuedResource, mockQueuedResource);
     }
   }
 
   @Test
   public void testDeleteTpuVm() {
     try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
       TpuClient mockTpuClient = mock(TpuClient.class);
+      OperationFuture mockFuture = mock(OperationFuture.class);
+
       mockedTpuClient.when(() -> TpuClient.create(any(TpuSettings.class)))
           .thenReturn(mockTpuClient);
-
-      OperationFuture mockFuture = mock(OperationFuture.class);
       when(mockTpuClient.deleteQueuedResourceAsync(any(DeleteQueuedResourceRequest.class)))
           .thenReturn(mockFuture);
+
       DeleteForceQueuedResource.deleteForceQueuedResource(PROJECT_ID, ZONE, QUEUED_RESOURCE_NAME);
       String output = bout.toString();
 

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

@@ -16,6 +16,7 @@
 
 package tpu;
 
+import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.mockStatic;
@@ -25,12 +26,15 @@
 
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.cloud.tpu.v2.DeleteNodeRequest;
+import com.google.cloud.tpu.v2.GetNodeRequest;
 import com.google.cloud.tpu.v2.Node;
-import com.google.cloud.tpu.v2.NodeName;
 import com.google.cloud.tpu.v2.TpuClient;
 import com.google.cloud.tpu.v2.TpuSettings;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.PrintStream;
 import java.util.concurrent.ExecutionException;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 import org.junit.runner.RunWith;
@@ -43,34 +47,47 @@ public class TpuVmIT {
   private static final String PROJECT_ID = "project-id";
   private static final String ZONE = "asia-east1-c";
   private static final String NODE_NAME = "test-tpu";
+  private static ByteArrayOutputStream bout;
+
+  @BeforeAll
+  public static void setUp() {
+    bout = new ByteArrayOutputStream();
+    System.setOut(new PrintStream(bout));
+  }
 
   @Test
   public void testGetTpuVm() throws IOException {
     try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
       Node mockNode = mock(Node.class);
-      mockedTpuClient.when(TpuClient::create).thenReturn(mock(TpuClient.class));
-      when(mock(TpuClient.class).getNode(any(NodeName.class))).thenReturn(mockNode);
+      TpuClient mockClient = mock(TpuClient.class);
       GetTpuVm mockGetTpuVm = mock(GetTpuVm.class);
 
-      GetTpuVm.getTpuVm(PROJECT_ID, ZONE, NODE_NAME);
+      mockedTpuClient.when(TpuClient::create).thenReturn(mockClient);
+      when(mockClient.getNode(any(GetNodeRequest.class))).thenReturn(mockNode);
+
+      Node returnedNode = GetTpuVm.getTpuVm(PROJECT_ID, ZONE, NODE_NAME);
 
       verify(mockGetTpuVm, times(1))
           .getTpuVm(PROJECT_ID, ZONE, NODE_NAME);
+      assertThat(returnedNode).isEqualTo(mockNode);
     }
   }
 
   @Test
   public void testDeleteTpuVm() throws IOException, ExecutionException, InterruptedException {
     try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
       TpuClient mockTpuClient = mock(TpuClient.class);
+      OperationFuture mockFuture = mock(OperationFuture.class);
+
       mockedTpuClient.when(() -> TpuClient.create(any(TpuSettings.class)))
           .thenReturn(mockTpuClient);
-
-      OperationFuture mockFuture = mock(OperationFuture.class);
       when(mockTpuClient.deleteNodeAsync(any(DeleteNodeRequest.class)))
           .thenReturn(mockFuture);
+
       DeleteTpuVm.deleteTpuVm(PROJECT_ID, ZONE, NODE_NAME);
+      String output = bout.toString();
 
+      assertThat(output).contains("TPU VM deleted");
       verify(mockTpuClient, times(1)).deleteNodeAsync(any(DeleteNodeRequest.class));
     }
   }