From 34e7cddd4ca729d388b4b3c1f45a7faf78b810cc Mon Sep 17 00:00:00 2001
From: Mike <45373284+munkhuushmgl@users.noreply.github.com>
Date: Fri, 28 Aug 2020 17:36:55 -0700
Subject: [PATCH] =?UTF-8?q?samples:=20refactor=20importDataset=20test=20se?=
 =?UTF-8?q?rvice=20randomly=20throws=20cancel=20e=E2=80=A6=20(#281)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* samples: refactor importDataset test service randomly throws cancel exception sometimes

* fixed the lint
---
 .../com/example/automl/ImportDatasetTest.java | 59 +++++++++++++++----
 1 file changed, 48 insertions(+), 11 deletions(-)

diff --git a/samples/snippets/src/test/java/com/example/automl/ImportDatasetTest.java b/samples/snippets/src/test/java/com/example/automl/ImportDatasetTest.java
index 9933a5d7d..e51652311 100644
--- a/samples/snippets/src/test/java/com/example/automl/ImportDatasetTest.java
+++ b/samples/snippets/src/test/java/com/example/automl/ImportDatasetTest.java
@@ -19,11 +19,20 @@
 import static com.google.common.truth.Truth.assertThat;
 import static junit.framework.TestCase.assertNotNull;
 
+import com.google.api.core.ApiFuture;
+import com.google.cloud.automl.v1beta1.AutoMlClient;
+import com.google.cloud.automl.v1beta1.CreateDatasetRequest;
+import com.google.cloud.automl.v1beta1.Dataset;
+import com.google.cloud.automl.v1beta1.LocationName;
+import com.google.cloud.automl.v1beta1.TextExtractionDatasetMetadata;
+import com.google.longrunning.Operation;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintStream;
 import java.util.UUID;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import org.junit.After;
 import org.junit.Before;
@@ -56,20 +65,34 @@ public static void checkRequirements() {
   }
 
   @Before
-  public void setUp() throws InterruptedException, ExecutionException, IOException {
-    bout = new ByteArrayOutputStream();
-    out = new PrintStream(bout);
-    System.setOut(out);
-
-    // Create a dataset that can be used for the import test
+  public void setUp()
+      throws IOException, InterruptedException, ExecutionException, TimeoutException {
+    // Create a fake dataset to be deleted
     // Create a random dataset name with a length of 32 characters (max allowed by AutoML)
     // To prevent name collisions when running tests in multiple java versions at once.
     // AutoML doesn't allow "-", but accepts "_"
     String datasetName =
         String.format("test_%s", UUID.randomUUID().toString().replace("-", "_").substring(0, 26));
-    LanguageEntityExtractionCreateDataset.createDataset(PROJECT_ID, datasetName);
-    String got = bout.toString();
-    datasetId = got.split("Dataset id: ")[1].split("\n")[0];
+    try (AutoMlClient client = AutoMlClient.create()) {
+
+      LocationName projectLocation = LocationName.of(PROJECT_ID, "us-central1");
+      TextExtractionDatasetMetadata metadata = TextExtractionDatasetMetadata.newBuilder().build();
+      Dataset dataset =
+          Dataset.newBuilder()
+              .setDisplayName(datasetName)
+              .setTextExtractionDatasetMetadata(metadata)
+              .build();
+
+      CreateDatasetRequest request =
+          CreateDatasetRequest.newBuilder()
+              .setParent(projectLocation.toString())
+              .setDataset(dataset)
+              .build();
+      ApiFuture<Dataset> future = client.createDatasetCallable().futureCall(request);
+      Dataset createdDataset = future.get(5, TimeUnit.MINUTES);
+      String[] names = createdDataset.getName().split("/");
+      datasetId = names[names.length - 1];
+    }
 
     bout = new ByteArrayOutputStream();
     out = new PrintStream(bout);
@@ -85,9 +108,23 @@ public void tearDown() throws InterruptedException, ExecutionException, IOExcept
 
   @Test
   public void testImportDataset()
-      throws IOException, ExecutionException, InterruptedException, TimeoutException {
-    ImportDataset.importDataset(PROJECT_ID, datasetId, BUCKET + "/entity-extraction/dataset.csv");
+      throws InterruptedException, ExecutionException, TimeoutException, IOException {
+
+    try {
+      ImportDataset.importDataset(PROJECT_ID, datasetId, BUCKET + "/entity-extraction/dataset.csv");
+    } catch (CancellationException ex) {
+      // capture operation ID from output and wait for that operation to be finished.
+      String fullOperationId = ex.getMessage().split("Operation name: ")[1].trim();
+      AutoMlClient client = AutoMlClient.create();
+      Operation importDatasetLro = client.getOperationsClient().getOperation(fullOperationId);
+      while (!importDatasetLro.getDone()) {
+        Thread.sleep(3000);
+      }
+      // retry the import.
+      ImportDataset.importDataset(PROJECT_ID, datasetId, BUCKET + "/entity-extraction/dataset.csv");
+    }
     String got = bout.toString();
+
     assertThat(got).contains("Dataset imported.");
   }
 }