diff --git a/automl/beta/src/test/java/com/example/automl/CancelOperationTest.java b/automl/beta/src/test/java/com/example/automl/CancelOperationTest.java index ea1f1dff638..916e80cbd70 100644 --- a/automl/beta/src/test/java/com/example/automl/CancelOperationTest.java +++ b/automl/beta/src/test/java/com/example/automl/CancelOperationTest.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.v1beta1.AutoMlClient; +import com.google.cloud.automl.v1beta1.DatasetName; +import com.google.cloud.automl.v1beta1.GcsSource; +import com.google.cloud.automl.v1beta1.InputConfig; +import com.google.cloud.automl.v1beta1.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.TimeoutException; import org.junit.After; import org.junit.Before; import org.junit.BeforeClass; @@ -32,16 +42,16 @@ @RunWith(JUnit4.class) public class CancelOperationTest { private static final String PROJECT_ID = System.getenv("AUTOML_PROJECT_ID"); - private static final String OPERATION_ID = "TRL6118473245806034944"; - private String operationFullNam = - String.format("projects/%s/locations/us-central1/operations/%s", PROJECT_ID, OPERATION_ID); + private static final String DATASET_ID = "TRL8242630754622767104"; + private static final String BUCKET = "gs://translate_data_exported/translation.csv"; + + private String operationId; + private String operationFullNam; private ByteArrayOutputStream bout; private PrintStream out; private static void requireEnvVar(String varName) { - assertNotNull( - System.getenv(varName), - "Environment variable '%s' is required to perform these tests.".format(varName)); + assertNotNull(System.getenv(varName), String.format(varName)); } @BeforeClass @@ -51,19 +61,49 @@ public static void checkRequirements() { } @Before - public void setUp() { + public void setUp() + throws InterruptedException, ExecutionException, TimeoutException, IOException { bout = new ByteArrayOutputStream(); out = new PrintStream(bout); System.setOut(out); + + // 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); + + // Get multiple Google Cloud Storage URIs to import data from + GcsSource gcsSource = + GcsSource.newBuilder().addAllInputUris(Arrays.asList(BUCKET.split(","))).build(); + + InputConfig inputConfig = InputConfig.newBuilder().setGcsSource(gcsSource).build(); + + // Start the import job + OperationFuture operation = + client.importDataAsync(datasetFullId, inputConfig); + + System.out.format("Operation name: %s%n", operation.getName()); + + operationFullNam = operation.getName(); + } } @After - public void tearDown() { + public void tearDown() throws IOException, InterruptedException { System.setOut(null); + + // delete the cancelled operation + try (AutoMlClient client = AutoMlClient.create()) { + // wait for the operation to be cancelled + while (!client.getOperationsClient().getOperation(operationFullNam).getDone()) { + Thread.sleep(1000); + } + client.getOperationsClient().deleteOperation(operationFullNam); + } } @Test - public void testGetOperationStatus() throws IOException { + public void testCancelOperation() throws IOException { CancelOperation.cancelOperation(operationFullNam); String got = bout.toString(); assertThat(got).contains("Operation cancelled");