modified tests to cancel job

GoogleCloudPlatform · Mar 19, 2020 · a1846b0 · a1846b0
1 parent fc181a9
commit a1846b0
Showing 1 changed file with 23 additions and 5 deletions.
diff --git a/dlp/src/test/java/dlp/snippets/InspectTests.java b/dlp/src/test/java/dlp/snippets/InspectTests.java
@@ -20,6 +20,8 @@
 import static org.hamcrest.core.StringContains.containsString;
 import static org.junit.Assert.assertNotNull;
 
+import com.google.cloud.dlp.v2.DlpServiceClient;
+import com.google.privacy.dlp.v2.CancelDlpJobRequest;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintStream;
@@ -103,8 +105,13 @@ public void testInspectGcsFile() throws InterruptedException, ExecutionException
     InspectGcsFile.inspectGcsFile(PROJECT_ID, GCS_PATH, TOPIC_ID, SUBSCRIPTION_ID);
 
     String output = bout.toString();
-    assertThat(output, containsString("Info type: PHONE_NUMBER"));
-    assertThat(output, containsString("Info type: EMAIL_ADDRESS"));
+    assertThat(output, containsString("Job created: "));
+
+    String jobId = output.split("Job created: ")[1].split("\n")[0];
+    CancelDlpJobRequest request = CancelDlpJobRequest.newBuilder().setName(jobId).build();
+    try (DlpServiceClient client = DlpServiceClient.create()) {
+      client.cancelDlpJob(request);
+    }
   }
 
   @Test
@@ -114,8 +121,13 @@ public void testInspectDatastoreEntity()
         PROJECT_ID, datastoreNamespace, datastoreKind, TOPIC_ID, SUBSCRIPTION_ID);
 
     String output = bout.toString();
-    assertThat(output, containsString("Info type: PHONE_NUMBER"));
-    assertThat(output, containsString("Info type: EMAIL_ADDRESS"));
+    assertThat(output, containsString("Job created: "));
+
+    String jobId = output.split("Job created: ")[1].split("\n")[0];
+    CancelDlpJobRequest request = CancelDlpJobRequest.newBuilder().setName(jobId).build();
+    try (DlpServiceClient client = DlpServiceClient.create()) {
+      client.cancelDlpJob(request);
+    }
   }
 
   @Test
@@ -125,6 +137,12 @@ public void testInspectBigQueryTable()
         PROJECT_ID, DATASET_ID, TABLE_ID, TOPIC_ID, SUBSCRIPTION_ID);
 
     String output = bout.toString();
-    assertThat(output, containsString("Info type: PHONE_NUMBER"));
+    assertThat(output, containsString("Job created: "));
+
+    String jobId = output.split("Job created: ")[1].split("\n")[0];
+    CancelDlpJobRequest request = CancelDlpJobRequest.newBuilder().setName(jobId).build();
+    try (DlpServiceClient client = DlpServiceClient.create()) {
+      client.cancelDlpJob(request);
+    }
   }
 }