Fixed test

Author: TetyanaYahodska

tpu/src/main/java/tpu/CreateSpotTpuVm.java

@@ -17,16 +17,12 @@
 package tpu;
 
 //[START tpu_vm_create_spot]
-import com.google.api.gax.longrunning.OperationTimedPollAlgorithm;
-import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.tpu.v2.CreateNodeRequest;
 import com.google.cloud.tpu.v2.Node;
 import com.google.cloud.tpu.v2.SchedulingConfig;
 import com.google.cloud.tpu.v2.TpuClient;
-import com.google.cloud.tpu.v2.TpuSettings;
 import java.io.IOException;
 import java.util.concurrent.ExecutionException;
-import org.threeten.bp.Duration;
 
 public class CreateSpotTpuVm {
   public static void main(String[] args)
@@ -37,7 +33,7 @@ public static void main(String[] args)
     // The zone in which to create the TPU.
     // For more information about supported TPU types for specific zones,
     // see https://cloud.google.com/tpu/docs/regions-zones
-    String zone = "europe-west4-a";
+    String zone = "us-central1-f";
     // The name for your TPU.
     String nodeName = "YOUR_TPY_NAME";
     // The accelerator type that specifies the version and size of the Cloud TPU you want to create.
@@ -55,27 +51,9 @@ public static void main(String[] args)
   public static Node createSpotTpuVm(
       String projectId, String zone, 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
-    TpuSettings.Builder clientSettings =
-        TpuSettings.newBuilder();
-    clientSettings
-        .createNodeOperationSettings()
-        .setPollingAlgorithm(
-            OperationTimedPollAlgorithm.create(
-                RetrySettings.newBuilder()
-                    .setInitialRetryDelay(Duration.ofMillis(5000L))
-                    .setRetryDelayMultiplier(1.5)
-                    .setMaxRetryDelay(Duration.ofMillis(45000L))
-                    .setInitialRpcTimeout(Duration.ZERO)
-                    .setRpcTimeoutMultiplier(1.0)
-                    .setMaxRpcTimeout(Duration.ZERO)
-                    .setTotalTimeout(Duration.ofHours(24L))
-                    .build()));
-
     // 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())) {
+    try (TpuClient tpuClient = TpuClient.create()) {
       String parent = String.format("projects/%s/locations/%s", projectId, zone);
       // TODO: Wait for update of library to change preemptible to spot=True
       SchedulingConfig schedulingConfig = SchedulingConfig.newBuilder()

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

@@ -148,8 +148,7 @@ public void testCreateSpotTpuVm() throws Exception {
       TpuClient mockTpuClient = mock(TpuClient.class);
       OperationFuture mockFuture = mock(OperationFuture.class);
 
-      mockedTpuClient.when(() -> TpuClient.create(any(TpuSettings.class)))
-          .thenReturn(mockTpuClient);
+      mockedTpuClient.when(TpuClient::create).thenReturn(mockTpuClient);
       when(mockTpuClient.createNodeAsync(any(CreateNodeRequest.class)))
           .thenReturn(mockFuture);
       when(mockFuture.get()).thenReturn(mockNode);