samples(test): refactor automl test (GoogleCloudPlatform#3588)

Fixes GoogleCloudPlatform#3550 Looks like in setup, ListOperations that had a very long list of operations to list which making a test to wait for 10-15min, Then, sometimes it fails with a internal error. - [X] Please **merge** this PR for me once it is approved.

Author: munkhuushmgl

automl/cloud-client/src/test/java/com/example/automl/GetOperationStatusTest.java

@@ -19,9 +19,19 @@
 import static com.google.common.truth.Truth.assertThat;
 import static junit.framework.TestCase.assertNotNull;
 
+import com.google.api.gax.longrunning.OperationFuture;
+import com.google.cloud.automl.v1.AutoMlClient;
+import com.google.cloud.automl.v1.DatasetName;
+import com.google.cloud.automl.v1.GcsSource;
+import com.google.cloud.automl.v1.InputConfig;
+import com.google.cloud.automl.v1.OperationMetadata;
+import com.google.protobuf.Empty;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintStream;
+import java.util.Arrays;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.BeforeClass;
@@ -32,7 +42,10 @@
 @RunWith(JUnit4.class)
 public class GetOperationStatusTest {
  private static final String PROJECT_ID = System.getenv("AUTOML_PROJECT_ID");
- private String operationId;
+ private static final String DATASET_ID = "TRL5394674636845744128";
+ private static final String BUCKET = "gs://translate_data_exported/translation.csv";
+
+ private String operationFullName;
  private ByteArrayOutputStream bout;
  private PrintStream out;
 
@@ -49,30 +62,48 @@ public static void checkRequirements() {
  }
 
  @Before
- public void setUp() throws IOException {
+ public void setUp() throws IOException, ExecutionException, InterruptedException {
  bout = new ByteArrayOutputStream();
  out = new PrintStream(bout);
  System.setOut(out);
 
- ListOperationStatus.listOperationStatus(PROJECT_ID);
- String got = bout.toString();
- operationId = got.split("\n")[1].split(":")[1].trim();
- assertThat(got).contains("Operation details:");
+ // start a export 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);
+
+ GcsSource gcsSource =
+ GcsSource.newBuilder().addAllInputUris(Arrays.asList(BUCKET.split(","))).build();
+
+ InputConfig inputConfig = InputConfig.newBuilder().setGcsSource(gcsSource).build();
+
+ // Start the import LRO job
+ OperationFuture<Empty, OperationMetadata> operation =
+ client.importDataAsync(datasetFullId, inputConfig);
+
+ operationFullName = operation.getName();
+ }
 
  bout = new ByteArrayOutputStream();
  out = new PrintStream(bout);
  System.setOut(out);
  }
 
- @After
- public void tearDown() {
- System.setOut(null);
- }
-
  @Test
  public void testGetOperationStatus() throws IOException {
- GetOperationStatus.getOperationStatus(operationId);
+ GetOperationStatus.getOperationStatus(operationFullName);
  String got = bout.toString();
  assertThat(got).contains("Operation details:");
  }
+
+ @After
+ public void tearDown() throws IOException, InterruptedException {
+ try (AutoMlClient client = AutoMlClient.create()) {
+ // terminate export data LRO.
+ client.getOperationsClient().cancelOperation(operationFullName);
+ client.getOperationsClient().awaitTermination(5, TimeUnit.SECONDS);
+ }
+
+ System.setOut(null);
+ }
 }