|
15 | 15 | import os
|
16 | 16 | import uuid
|
17 | 17 |
|
18 |
| -from gcp_devrel.testing import eventually_consistent |
19 |
| -from gcp_devrel.testing.flaky import flaky |
20 | 18 | import google.api_core.exceptions
|
21 | 19 | import google.cloud.bigquery
|
22 | 20 | import google.cloud.datastore
|
23 | 21 | import google.cloud.dlp_v2
|
24 | 22 | import google.cloud.exceptions
|
25 | 23 | import google.cloud.pubsub
|
26 | 24 | import google.cloud.storage
|
27 |
| - |
28 | 25 | import pytest
|
| 26 | + |
29 | 27 | import inspect_content
|
30 | 28 |
|
| 29 | + |
31 | 30 | UNIQUE_STRING = str(uuid.uuid4()).split("-")[0]
|
32 | 31 |
|
33 | 32 | GCLOUD_PROJECT = os.getenv("GCLOUD_PROJECT")
|
@@ -95,7 +94,8 @@ def subscription_id(topic_id):
|
95 | 94 | # Subscribes to a topic.
|
96 | 95 | subscriber = google.cloud.pubsub.SubscriberClient()
|
97 | 96 | topic_path = subscriber.topic_path(GCLOUD_PROJECT, topic_id)
|
98 |
| - subscription_path = subscriber.subscription_path(GCLOUD_PROJECT, SUBSCRIPTION_ID) |
| 97 | + subscription_path = subscriber.subscription_path( |
| 98 | + GCLOUD_PROJECT, SUBSCRIPTION_ID) |
99 | 99 | try:
|
100 | 100 | subscriber.create_subscription(subscription_path, topic_path)
|
101 | 101 | except google.api_core.exceptions.AlreadyExists:
|
@@ -289,157 +289,160 @@ def test_inspect_image_file(capsys):
|
289 | 289 | assert "Info type: PHONE_NUMBER" in out
|
290 | 290 |
|
291 | 291 |
|
| 292 | +def cancel_operation(out): |
| 293 | + if "Inspection operation started" in out: |
| 294 | + # Cancel the operation |
| 295 | + operation_id = out.split( |
| 296 | + "Inspection operation started: ")[1].split("\n")[0] |
| 297 | + client = google.cloud.dlp_v2.DlpServiceClient() |
| 298 | + client.cancel_dlp_job(operation_id) |
| 299 | + |
| 300 | + |
292 | 301 | def test_inspect_gcs_file(bucket, topic_id, subscription_id, capsys):
|
293 |
| - inspect_content.inspect_gcs_file( |
294 |
| - GCLOUD_PROJECT, |
295 |
| - bucket.name, |
296 |
| - "test.txt", |
297 |
| - topic_id, |
298 |
| - subscription_id, |
299 |
| - ["EMAIL_ADDRESS", "PHONE_NUMBER"], |
300 |
| - timeout=420, |
301 |
| - ) |
| 302 | + try: |
| 303 | + inspect_content.inspect_gcs_file( |
| 304 | + GCLOUD_PROJECT, |
| 305 | + bucket.name, |
| 306 | + "test.txt", |
| 307 | + topic_id, |
| 308 | + subscription_id, |
| 309 | + ["EMAIL_ADDRESS", "PHONE_NUMBER"], |
| 310 | + timeout=1 |
| 311 | + ) |
302 | 312 |
|
303 |
| - out, _ = capsys.readouterr() |
304 |
| - assert "Inspection operation started" in out |
305 |
| - # Cancel the operation |
306 |
| - operation_id = out.split("Inspection operation started: ")[1].split("\n")[0] |
307 |
| - print(operation_id) |
308 |
| - client = google.cloud.dlp_v2.DlpServiceClient() |
309 |
| - client.cancel_dlp_job(operation_id) |
| 313 | + out, _ = capsys.readouterr() |
| 314 | + assert "Inspection operation started" in out |
| 315 | + finally: |
| 316 | + cancel_operation(out) |
310 | 317 |
|
311 | 318 |
|
312 | 319 | def test_inspect_gcs_file_with_custom_info_types(
|
313 |
| - bucket, topic_id, subscription_id, capsys |
314 |
| -): |
315 |
| - dictionaries = [ "[email protected]"] |
316 |
| - regexes = ["\\(\\d{3}\\) \\d{3}-\\d{4}"] |
| 320 | + bucket, topic_id, subscription_id, capsys): |
| 321 | + try: |
| 322 | + dictionaries = [ "[email protected]"] |
| 323 | + regexes = ["\\(\\d{3}\\) \\d{3}-\\d{4}"] |
317 | 324 |
|
318 |
| - inspect_content.inspect_gcs_file( |
319 |
| - GCLOUD_PROJECT, |
320 |
| - bucket.name, |
321 |
| - "test.txt", |
322 |
| - topic_id, |
323 |
| - subscription_id, |
324 |
| - [], |
325 |
| - custom_dictionaries=dictionaries, |
326 |
| - custom_regexes=regexes, |
327 |
| - timeout=420, |
328 |
| - ) |
| 325 | + inspect_content.inspect_gcs_file( |
| 326 | + GCLOUD_PROJECT, |
| 327 | + bucket.name, |
| 328 | + "test.txt", |
| 329 | + topic_id, |
| 330 | + subscription_id, |
| 331 | + [], |
| 332 | + custom_dictionaries=dictionaries, |
| 333 | + custom_regexes=regexes, |
| 334 | + timeout=1) |
329 | 335 |
|
330 |
| - out, _ = capsys.readouterr() |
| 336 | + out, _ = capsys.readouterr() |
331 | 337 |
|
332 |
| - assert "Inspection operation started" in out |
333 |
| - # Cancel the operation |
334 |
| - operation_id = out.split("Inspection operation started: ")[1].split("\n")[0] |
335 |
| - print(operation_id) |
336 |
| - client = google.cloud.dlp_v2.DlpServiceClient() |
337 |
| - client.cancel_dlp_job(operation_id) |
| 338 | + assert "Inspection operation started" in out |
| 339 | + finally: |
| 340 | + cancel_operation(out) |
338 | 341 |
|
339 | 342 |
|
340 |
| -def test_inspect_gcs_file_no_results(bucket, topic_id, subscription_id, capsys): |
341 |
| - inspect_content.inspect_gcs_file( |
342 |
| - GCLOUD_PROJECT, |
343 |
| - bucket.name, |
344 |
| - "harmless.txt", |
345 |
| - topic_id, |
346 |
| - subscription_id, |
347 |
| - ["EMAIL_ADDRESS", "PHONE_NUMBER"], |
348 |
| - timeout=420, |
349 |
| - ) |
| 343 | +def test_inspect_gcs_file_no_results( |
| 344 | + bucket, topic_id, subscription_id, capsys): |
| 345 | + try: |
| 346 | + inspect_content.inspect_gcs_file( |
| 347 | + GCLOUD_PROJECT, |
| 348 | + bucket.name, |
| 349 | + "harmless.txt", |
| 350 | + topic_id, |
| 351 | + subscription_id, |
| 352 | + ["EMAIL_ADDRESS", "PHONE_NUMBER"], |
| 353 | + timeout=1) |
350 | 354 |
|
351 |
| - out, _ = capsys.readouterr() |
| 355 | + out, _ = capsys.readouterr() |
352 | 356 |
|
353 |
| - assert "Inspection operation started" in out |
354 |
| - # Cancel the operation |
355 |
| - operation_id = out.split("Inspection operation started: ")[1].split("\n")[0] |
356 |
| - print(operation_id) |
357 |
| - client = google.cloud.dlp_v2.DlpServiceClient() |
358 |
| - client.cancel_dlp_job(operation_id) |
| 357 | + assert "Inspection operation started" in out |
| 358 | + finally: |
| 359 | + cancel_operation(out) |
359 | 360 |
|
360 | 361 |
|
361 |
| -@pytest.mark.skip(reason="nondeterministically failing") |
362 | 362 | def test_inspect_gcs_image_file(bucket, topic_id, subscription_id, capsys):
|
363 |
| - inspect_content.inspect_gcs_file( |
364 |
| - GCLOUD_PROJECT, |
365 |
| - bucket.name, |
366 |
| - "test.png", |
367 |
| - topic_id, |
368 |
| - subscription_id, |
369 |
| - ["EMAIL_ADDRESS", "PHONE_NUMBER"], |
370 |
| - ) |
| 363 | + try: |
| 364 | + inspect_content.inspect_gcs_file( |
| 365 | + GCLOUD_PROJECT, |
| 366 | + bucket.name, |
| 367 | + "test.png", |
| 368 | + topic_id, |
| 369 | + subscription_id, |
| 370 | + ["EMAIL_ADDRESS", "PHONE_NUMBER"], |
| 371 | + timeout=1) |
371 | 372 |
|
372 |
| - out, _ = capsys.readouterr() |
373 |
| - assert "Info type: EMAIL_ADDRESS" in out |
| 373 | + out, _ = capsys.readouterr() |
| 374 | + assert "Inspection operation started" in out |
| 375 | + finally: |
| 376 | + cancel_operation(out) |
374 | 377 |
|
375 | 378 |
|
376 | 379 | def test_inspect_gcs_multiple_files(bucket, topic_id, subscription_id, capsys):
|
377 |
| - inspect_content.inspect_gcs_file( |
378 |
| - GCLOUD_PROJECT, |
379 |
| - bucket.name, |
380 |
| - "*", |
381 |
| - topic_id, |
382 |
| - subscription_id, |
383 |
| - ["EMAIL_ADDRESS", "PHONE_NUMBER"], |
384 |
| - ) |
| 380 | + try: |
| 381 | + inspect_content.inspect_gcs_file( |
| 382 | + GCLOUD_PROJECT, |
| 383 | + bucket.name, |
| 384 | + "*", |
| 385 | + topic_id, |
| 386 | + subscription_id, |
| 387 | + ["EMAIL_ADDRESS", "PHONE_NUMBER"], |
| 388 | + timeout=1) |
385 | 389 |
|
386 |
| - out, _ = capsys.readouterr() |
| 390 | + out, _ = capsys.readouterr() |
387 | 391 |
|
388 |
| - assert "Inspection operation started" in out |
389 |
| - # Cancel the operation |
390 |
| - operation_id = out.split("Inspection operation started: ")[1].split("\n")[0] |
391 |
| - print(operation_id) |
392 |
| - client = google.cloud.dlp_v2.DlpServiceClient() |
393 |
| - client.cancel_dlp_job(operation_id) |
| 392 | + assert "Inspection operation started" in out |
| 393 | + finally: |
| 394 | + cancel_operation(out) |
394 | 395 |
|
395 | 396 |
|
396 |
| -@flaky |
397 |
| -def test_inspect_datastore(datastore_project, topic_id, subscription_id, capsys): |
398 |
| - @eventually_consistent.call |
399 |
| - def _(): |
| 397 | +def test_inspect_datastore( |
| 398 | + datastore_project, topic_id, subscription_id, capsys): |
| 399 | + try: |
400 | 400 | inspect_content.inspect_datastore(
|
401 | 401 | GCLOUD_PROJECT,
|
402 | 402 | datastore_project,
|
403 | 403 | DATASTORE_KIND,
|
404 | 404 | topic_id,
|
405 | 405 | subscription_id,
|
406 | 406 | ["FIRST_NAME", "EMAIL_ADDRESS", "PHONE_NUMBER"],
|
407 |
| - ) |
| 407 | + timeout=1) |
408 | 408 |
|
409 | 409 | out, _ = capsys.readouterr()
|
410 |
| - assert "Info type: EMAIL_ADDRESS" in out |
| 410 | + assert "Inspection operation started" in out |
| 411 | + finally: |
| 412 | + cancel_operation(out) |
411 | 413 |
|
412 | 414 |
|
413 |
| -@flaky |
414 | 415 | def test_inspect_datastore_no_results(
|
415 |
| - datastore_project, topic_id, subscription_id, capsys |
416 |
| -): |
417 |
| - @eventually_consistent.call |
418 |
| - def _(): |
| 416 | + datastore_project, topic_id, subscription_id, capsys): |
| 417 | + try: |
419 | 418 | inspect_content.inspect_datastore(
|
420 | 419 | GCLOUD_PROJECT,
|
421 | 420 | datastore_project,
|
422 | 421 | DATASTORE_KIND,
|
423 | 422 | topic_id,
|
424 | 423 | subscription_id,
|
425 | 424 | ["PHONE_NUMBER"],
|
426 |
| - ) |
| 425 | + timeout=1) |
427 | 426 |
|
428 | 427 | out, _ = capsys.readouterr()
|
429 |
| - assert "No findings" in out |
| 428 | + assert "Inspection operation started" in out |
| 429 | + finally: |
| 430 | + cancel_operation(out) |
430 | 431 |
|
431 | 432 |
|
432 |
| -@pytest.mark.skip(reason="unknown issue") |
433 | 433 | def test_inspect_bigquery(bigquery_project, topic_id, subscription_id, capsys):
|
434 |
| - inspect_content.inspect_bigquery( |
435 |
| - GCLOUD_PROJECT, |
436 |
| - bigquery_project, |
437 |
| - BIGQUERY_DATASET_ID, |
438 |
| - BIGQUERY_TABLE_ID, |
439 |
| - topic_id, |
440 |
| - subscription_id, |
441 |
| - ["FIRST_NAME", "EMAIL_ADDRESS", "PHONE_NUMBER"], |
442 |
| - ) |
| 434 | + try: |
| 435 | + inspect_content.inspect_bigquery( |
| 436 | + GCLOUD_PROJECT, |
| 437 | + bigquery_project, |
| 438 | + BIGQUERY_DATASET_ID, |
| 439 | + BIGQUERY_TABLE_ID, |
| 440 | + topic_id, |
| 441 | + subscription_id, |
| 442 | + ["FIRST_NAME", "EMAIL_ADDRESS", "PHONE_NUMBER"], |
| 443 | + timeout=1) |
443 | 444 |
|
444 |
| - out, _ = capsys.readouterr() |
445 |
| - assert "Info type: FIRST_NAME" in out |
| 445 | + out, _ = capsys.readouterr() |
| 446 | + assert "Inspection operation started" in out |
| 447 | + finally: |
| 448 | + cancel_operation(out) |
0 commit comments