samples: fixed flaky test on getOper sample (GoogleCloudPlatform#3675)

munkhuushmgl · web-flow · commit d1514d33dcae · 2020-09-08T21:52:02.000Z
- [X] Please **merge** this PR for me once it is approved.
diff --git a/automl/cloud-client/src/main/java/com/example/automl/GetOperationStatus.java b/automl/cloud-client/src/main/java/com/example/automl/GetOperationStatus.java
@@ -19,6 +19,7 @@
 // [START automl_get_operation_status]
 import com.google.cloud.automl.v1.AutoMlClient;
 import com.google.longrunning.Operation;
+import io.grpc.StatusRuntimeException;
 import java.io.IOException;
 
 class GetOperationStatus {
@@ -30,7 +31,8 @@ static void getOperationStatus() throws IOException {
   }
 
   // Get the status of an operation
-  static void getOperationStatus(String operationFullId) throws IOException {
+  static void getOperationStatus(String operationFullId)
+      throws IOException, StatusRuntimeException {
     // Initialize client that will be used to send requests. This client only needs to be created
     // once, and can be reused for multiple requests. After completing all of your requests, call
     // the "close" method on the client to safely clean up any remaining background resources.
diff --git a/automl/cloud-client/src/test/java/com/example/automl/GetOperationStatusTest.java b/automl/cloud-client/src/test/java/com/example/automl/GetOperationStatusTest.java
@@ -67,7 +67,7 @@ public void setUp() throws IOException, ExecutionException, InterruptedException
     out = new PrintStream(bout);
     System.setOut(out);
 
-    // start a export data into dataset and cancel it before it finishes.
+    // start a import data into dataset and cancel it before it finishes.
     try (AutoMlClient client = AutoMlClient.create()) {
       // Get the complete path of the dataset.
       DatasetName datasetFullId = DatasetName.of(PROJECT_ID, "us-central1", DATASET_ID);
@@ -99,9 +99,11 @@ public void testGetOperationStatus() throws IOException {
   @After
   public void tearDown() throws IOException, InterruptedException {
     try (AutoMlClient client = AutoMlClient.create()) {
-      // terminate export data LRO.
+      // terminate import data LRO.
       client.getOperationsClient().cancelOperation(operationFullName);
-      client.getOperationsClient().awaitTermination(5, TimeUnit.SECONDS);
+      while (!client.getOperationsClient().getOperation(operationFullName).getDone()) {
+        TimeUnit.SECONDS.sleep(5);
+      }
     }
 
     System.setOut(null);
