Skip to content

Commit ef3511d

Browse files
[DLP] Mock flaky unit test cases (GoogleCloudPlatform#8430)
* mocked unit test cases * mocked unit test cases
1 parent ef6d07b commit ef3511d

2 files changed

Lines changed: 349 additions & 159 deletions

File tree

‎dlp/snippets/src/test/java/dlp/snippets/InspectTests.java‎

Lines changed: 108 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -337,18 +337,43 @@ public void testRedactImageListedInfoTypes() throws Exception {
337337

338338
@Test
339339
public void testInspectGcsFile() throws Exception {
340-
InspectGcsFile.inspectGcsFile(
341-
PROJECT_ID, GCS_PATH, topicName.getTopic(), subscriptionName.getSubscription());
340+
DlpServiceClient dlpServiceClient = mock(DlpServiceClient.class);
341+
doneMock = mock(SettableApiFuture.class);
342+
try (MockedStatic<DlpServiceClient> mockedStatic = Mockito.mockStatic(DlpServiceClient.class)) {
343+
mockedStatic.when(() -> DlpServiceClient.create()).thenReturn(dlpServiceClient);
344+
try (MockedStatic<SettableApiFuture> mockedStatic1 =
345+
Mockito.mockStatic(SettableApiFuture.class)) {
346+
mockedStatic1.when(() -> SettableApiFuture.create()).thenReturn(doneMock);
342347

343-
String output = bout.toString();
344-
assertThat(output).contains("Job status: DONE");
345-
String jobName = Arrays.stream(output.split("\n"))
346-
.filter(line -> line.contains("Job name:"))
347-
.findFirst()
348-
.get();
349-
jobName = jobName.split(":")[1].trim();
350-
try (DlpServiceClient dlp = DlpServiceClient.create()) {
351-
dlp.deleteDlpJob(jobName);
348+
InfoTypeStats infoTypeStats =
349+
InfoTypeStats.newBuilder()
350+
.setInfoType(InfoType.newBuilder().setName("EMAIL_ADDRESS").build())
351+
.setCount(1)
352+
.build();
353+
DlpJob dlpJob =
354+
DlpJob.newBuilder()
355+
.setName("projects/project_id/locations/global/dlpJobs/job_id")
356+
.setState(DlpJob.JobState.DONE)
357+
.setInspectDetails(
358+
InspectDataSourceDetails.newBuilder()
359+
.setResult(
360+
InspectDataSourceDetails.Result.newBuilder()
361+
.addInfoTypeStats(infoTypeStats)
362+
.build()))
363+
.build();
364+
when(doneMock.get(15, TimeUnit.MINUTES)).thenReturn(true);
365+
when(dlpServiceClient.createDlpJob(any(CreateDlpJobRequest.class))).thenReturn(dlpJob);
366+
when(dlpServiceClient.getDlpJob((GetDlpJobRequest) any())).thenReturn(dlpJob);
367+
InspectGcsFile.inspectGcsFile(
368+
"project_id", "gs://bucket_name/test.txt", "topic_id", "subscription_id");
369+
String output = bout.toString();
370+
assertThat(output).contains("Job status: DONE");
371+
assertThat(output)
372+
.contains("Job name: projects/project_id/locations/global/dlpJobs/job_id");
373+
assertThat(output).contains("Info type: EMAIL_ADDRESS");
374+
verify(dlpServiceClient, times(1)).createDlpJob(any(CreateDlpJobRequest.class));
375+
verify(dlpServiceClient, times(1)).getDlpJob(any(GetDlpJobRequest.class));
376+
}
352377
}
353378
}
354379

@@ -396,39 +421,85 @@ public void testInspectGcsFileWithSampling() throws Exception {
396421

397422
@Test
398423
public void testInspectDatastoreEntity() throws Exception {
399-
InspectDatastoreEntity.insepctDatastoreEntity(
400-
PROJECT_ID,
401-
datastoreNamespace,
402-
datastoreKind,
403-
topicName.getTopic(),
404-
subscriptionName.getSubscription());
424+
DlpServiceClient dlpServiceClient = mock(DlpServiceClient.class);
425+
doneMock = mock(SettableApiFuture.class);
426+
try (MockedStatic<DlpServiceClient> mockedStatic = Mockito.mockStatic(DlpServiceClient.class)) {
427+
mockedStatic.when(() -> DlpServiceClient.create()).thenReturn(dlpServiceClient);
428+
try (MockedStatic<SettableApiFuture> mockedStatic1 =
429+
Mockito.mockStatic(SettableApiFuture.class)) {
430+
mockedStatic1.when(() -> SettableApiFuture.create()).thenReturn(doneMock);
405431

406-
String output = bout.toString();
407-
assertThat(output).contains("Job status: DONE");
408-
String jobName = Arrays.stream(output.split("\n"))
409-
.filter(line -> line.contains("Job name:"))
410-
.findFirst()
411-
.get();
412-
jobName = jobName.split(":")[1].trim();
413-
try (DlpServiceClient dlp = DlpServiceClient.create()) {
414-
dlp.deleteDlpJob(jobName);
432+
InfoTypeStats infoTypeStats =
433+
InfoTypeStats.newBuilder()
434+
.setInfoType(InfoType.newBuilder().setName("EMAIL_ADDRESS").build())
435+
.setCount(1)
436+
.build();
437+
DlpJob dlpJob =
438+
DlpJob.newBuilder()
439+
.setName("projects/project_id/locations/global/dlpJobs/job_id")
440+
.setState(DlpJob.JobState.DONE)
441+
.setInspectDetails(
442+
InspectDataSourceDetails.newBuilder()
443+
.setResult(
444+
InspectDataSourceDetails.Result.newBuilder()
445+
.addInfoTypeStats(infoTypeStats)
446+
.build()))
447+
.build();
448+
when(doneMock.get(15, TimeUnit.MINUTES)).thenReturn(true);
449+
when(dlpServiceClient.createDlpJob(any(CreateDlpJobRequest.class))).thenReturn(dlpJob);
450+
when(dlpServiceClient.getDlpJob((GetDlpJobRequest) any())).thenReturn(dlpJob);
451+
InspectDatastoreEntity.insepctDatastoreEntity("project_id", "datastore_namespace_test",
452+
"datastore_kind_test", "topic_id", "subscription_id");
453+
String output = bout.toString();
454+
assertThat(output).contains("Job status: DONE");
455+
assertThat(output)
456+
.contains("Job name: projects/project_id/locations/global/dlpJobs/job_id");
457+
assertThat(output).contains("Info type: EMAIL_ADDRESS");
458+
verify(dlpServiceClient, times(1)).createDlpJob(any(CreateDlpJobRequest.class));
459+
verify(dlpServiceClient, times(1)).getDlpJob(any(GetDlpJobRequest.class));
460+
}
415461
}
416462
}
417463

418464
@Test
419465
public void testInspectBigQueryTable() throws Exception {
420-
InspectBigQueryTable.inspectBigQueryTable(
421-
PROJECT_ID, DATASET_ID, TABLE_ID, topicName.getTopic(), subscriptionName.getSubscription());
466+
DlpServiceClient dlpServiceClient = mock(DlpServiceClient.class);
467+
doneMock = mock(SettableApiFuture.class);
468+
try (MockedStatic<DlpServiceClient> mockedStatic = Mockito.mockStatic(DlpServiceClient.class)) {
469+
mockedStatic.when(() -> DlpServiceClient.create()).thenReturn(dlpServiceClient);
470+
try (MockedStatic<SettableApiFuture> mockedStatic1 =
471+
Mockito.mockStatic(SettableApiFuture.class)) {
472+
mockedStatic1.when(() -> SettableApiFuture.create()).thenReturn(doneMock);
422473

423-
String output = bout.toString();
424-
assertThat(output).contains("Job status: DONE");
425-
String jobName = Arrays.stream(output.split("\n"))
426-
.filter(line -> line.contains("Job name:"))
427-
.findFirst()
428-
.get();
429-
jobName = jobName.split(":")[1].trim();
430-
try (DlpServiceClient dlp = DlpServiceClient.create()) {
431-
dlp.deleteDlpJob(jobName);
474+
InfoTypeStats infoTypeStats =
475+
InfoTypeStats.newBuilder()
476+
.setInfoType(InfoType.newBuilder().setName("EMAIL_ADDRESS").build())
477+
.setCount(1)
478+
.build();
479+
DlpJob dlpJob =
480+
DlpJob.newBuilder()
481+
.setName("projects/project_id/locations/global/dlpJobs/job_id")
482+
.setState(DlpJob.JobState.DONE)
483+
.setInspectDetails(
484+
InspectDataSourceDetails.newBuilder()
485+
.setResult(
486+
InspectDataSourceDetails.Result.newBuilder()
487+
.addInfoTypeStats(infoTypeStats)
488+
.build()))
489+
.build();
490+
when(doneMock.get(15, TimeUnit.MINUTES)).thenReturn(true);
491+
when(dlpServiceClient.createDlpJob(any(CreateDlpJobRequest.class))).thenReturn(dlpJob);
492+
when(dlpServiceClient.getDlpJob((GetDlpJobRequest) any())).thenReturn(dlpJob);
493+
InspectBigQueryTable.inspectBigQueryTable("bigquery-public-data", "usa_names",
494+
"usa_1910_current", "topic_id", "subscription_id");
495+
String output = bout.toString();
496+
assertThat(output).contains("Job status: DONE");
497+
assertThat(output)
498+
.contains("Job name: projects/project_id/locations/global/dlpJobs/job_id");
499+
assertThat(output).contains("Info type: EMAIL_ADDRESS");
500+
verify(dlpServiceClient, times(1)).createDlpJob(any(CreateDlpJobRequest.class));
501+
verify(dlpServiceClient, times(1)).getDlpJob(any(GetDlpJobRequest.class));
502+
}
432503
}
433504
}
434505

0 commit comments

Comments
 (0)