Fixed test

GoogleCloudPlatform · Dec 19, 2024 · b58fc0f · b58fc0f
1 parent 8e98108
commit b58fc0f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/tpu/src/main/java/tpu/CreateTimeBoundQueuedResource.java b/tpu/src/main/java/tpu/CreateTimeBoundQueuedResource.java
@@ -144,8 +144,6 @@ public static QueuedResource createTimeBoundQueuedResource(
               .setQueuedResourceId(queuedResourceName)
               .build();
 
-      // You can wait until TPU Node is READY,
-      // and check its status using getTpuVm() from "tpu_vm_get" sample.
       return tpuClient.createQueuedResourceAsync(request).get();
     }
   }

diff --git a/tpu/src/test/java/tpu/QueuedResourceIT.java b/tpu/src/test/java/tpu/QueuedResourceIT.java
@@ -140,7 +140,9 @@ public void testCreateQueuedResourceWithStartupScript() throws Exception {
   @Test
   public void testCreateTimeBoundQueuedResource() throws Exception {
     try (MockedStatic<TpuClient> mockedTpuClient = mockStatic(TpuClient.class)) {
-      QueuedResource mockQueuedResource = mock(QueuedResource.class);
+      QueuedResource mockQueuedResource =  QueuedResource.newBuilder()
+              .setName("QueuedResourceName")
+              .build();
       TpuClient mockTpuClient = mock(TpuClient.class);
       OperationFuture mockFuture = mock(OperationFuture.class);
 
@@ -157,7 +159,7 @@ public void testCreateTimeBoundQueuedResource() throws Exception {
       verify(mockTpuClient, times(1))
           .createQueuedResourceAsync(any(CreateQueuedResourceRequest.class));
       verify(mockFuture, times(1)).get();
-      assertEquals(returnedQueuedResource, mockQueuedResource);
+      assertEquals(returnedQueuedResource.getName(), mockQueuedResource.getName());
     }
   }
 }