GoogleCloudPlatform
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
@@ -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);
+  }
 }