fix DLP region tags, and add @flaky to pub/sub sample tests (#1418)

Author: andrewsg

dlp/deid.py

@@ -19,6 +19,7 @@
 import argparse
 
 
+# [START dlp_deidentify_masking]
 def deidentify_with_mask(project, string, masking_character=None,
                          number_to_mask=0):
     """Uses the Data Loss Prevention API to deidentify sensitive data in a
@@ -68,8 +69,10 @@ def deidentify_with_mask(project, string, masking_character=None,
 
     # Print out the results.
     print(response.item.value)
+# [END dlp_deidentify_masking]
 
 
+# [START dlp_deidentify_fpe]
 def deidentify_with_fpe(project, string, alphabet=None,
                         surrogate_type=None, key_name=None, wrapped_key=None):
     """Uses the Data Loss Prevention API to deidentify sensitive data in a
@@ -147,8 +150,10 @@ def deidentify_with_fpe(project, string, alphabet=None,
 
     # Print results
     print(response.item.value)
+# [END dlp_deidentify_fpe]
 
 
+# [START dlp_reidentify_fpe]
 def reidentify_with_fpe(project, string, alphabet=None,
                         surrogate_type=None, key_name=None, wrapped_key=None):
     """Uses the Data Loss Prevention API to reidentify sensitive data in a
@@ -232,8 +237,10 @@ def reidentify_with_fpe(project, string, alphabet=None,
 
     # Print results
     print(response.item.value)
+# [END dlp_reidentify_fpe]
 
 
+# [START dlp_deidentify_date_shift]
 def deidentify_with_date_shift(project, input_csv_file=None,
                                output_csv_file=None, date_fields=None,
                                lower_bound_days=None, upper_bound_days=None,
@@ -384,6 +391,7 @@ def write_data(data):
     # Print status
     print('Successfully saved date-shift output to {}'.format(
                 output_csv_file))
+# [END dlp_deidentify_date_shift]
 
 
 if __name__ == '__main__':

dlp/inspect_content_test.py

@@ -222,6 +222,7 @@ def test_inspect_image_file(capsys):
     assert 'Info type: PHONE_NUMBER' in out
 
 
+@flaky
 def test_inspect_gcs_file(bucket, topic_id, subscription_id, capsys):
     inspect_content.inspect_gcs_file(
         GCLOUD_PROJECT,
@@ -235,6 +236,7 @@ def test_inspect_gcs_file(bucket, topic_id, subscription_id, capsys):
     assert 'Info type: EMAIL_ADDRESS' in out
 
 
+@flaky
 def test_inspect_gcs_file_no_results(
         bucket, topic_id, subscription_id, capsys):
     inspect_content.inspect_gcs_file(
@@ -263,6 +265,7 @@ def test_inspect_gcs_image_file(bucket, topic_id, subscription_id, capsys):
     assert 'Info type: EMAIL_ADDRESS' in out
 
 
+@flaky
 def test_inspect_gcs_multiple_files(bucket, topic_id, subscription_id, capsys):
     inspect_content.inspect_gcs_file(
         GCLOUD_PROJECT,
@@ -294,6 +297,7 @@ def _():
         assert 'Info type: EMAIL_ADDRESS' in out
 
 
+@flaky
 def test_inspect_datastore_no_results(
         datastore_project, topic_id, subscription_id, capsys):
     inspect_content.inspect_datastore(

dlp/jobs.py

@@ -19,6 +19,7 @@
 import argparse
 
 
+# [START dlp_list_jobs]
 def list_dlp_jobs(project, filter_string=None, job_type=None):
     """Uses the Data Loss Prevention API to lists DLP jobs that match the
         specified filter in the request.
@@ -83,8 +84,10 @@ def list_dlp_jobs(project, filter_string=None, job_type=None):
     # Iterate over results.
     for job in response:
         print('Job: %s; status: %s' % (job.name, job.JobState.Name(job.state)))
+# [END dlp_list_jobs]
 
 
+# [START dlp_delete_job]
 def delete_dlp_job(project, job_name):
     """Uses the Data Loss Prevention API to delete a long-running DLP job.
     Args:
@@ -108,6 +111,7 @@ def delete_dlp_job(project, job_name):
     dlp.delete_dlp_job(name)
 
     print('Successfully deleted %s' % job_name)
+# [END dlp_delete_job]
 
 
 if __name__ == '__main__':

dlp/metadata.py

@@ -20,7 +20,7 @@
 import argparse
 
 
-# [START list_info_types]
+# [START dlp_list_info_types]
 def list_info_types(language_code=None, result_filter=None):
     """List types of sensitive information within a category.
     Args:
@@ -44,7 +44,7 @@ def list_info_types(language_code=None, result_filter=None):
     for info_type in response.info_types:
         print('{name}: {display_name}'.format(
             name=info_type.name, display_name=info_type.display_name))
-# [END list_info_types]
+# [END dlp_list_info_types]
 
 
 if __name__ == '__main__':

dlp/quickstart.py

@@ -21,7 +21,7 @@
 def quickstart():
     """Demonstrates use of the Data Loss Prevention API client library."""
 
-    # [START quickstart]
+    # [START dlp_quickstart]
     # Import the client library
     import google.cloud.dlp
 
@@ -80,7 +80,7 @@ def quickstart():
             print('Likelihood: {}'.format(likelihood))
     else:
         print('No findings.')
-    # [END quickstart]
+    # [END dlp_quickstart]
 
 
 if __name__ == '__main__':

dlp/risk.py

@@ -19,6 +19,7 @@
 import argparse
 
 
+# [START dlp_numerical_stats]
 def numerical_risk_analysis(project, table_project_id, dataset_id, table_id,
                             column_name, topic_id, subscription_id,
                             timeout=300):
@@ -133,8 +134,10 @@ def callback(message):
     if not finished:
         print('No event received before the timeout. Please verify that the '
               'subscription provided is subscribed to the topic provided.')
+# [END dlp_numerical_stats]
 
 
+# [START dlp_categorical_stats]
 def categorical_risk_analysis(project, table_project_id, dataset_id, table_id,
                               column_name, topic_id, subscription_id,
                               timeout=300):
@@ -253,8 +256,10 @@ def callback(message):
     if not finished:
         print('No event received before the timeout. Please verify that the '
               'subscription provided is subscribed to the topic provided.')
+# [END dlp_categorical_stats]
 
 
+# [START dlp_k_anonymity]
 def k_anonymity_analysis(project, table_project_id, dataset_id, table_id,
                          topic_id, subscription_id, quasi_ids, timeout=300):
     """Uses the Data Loss Prevention API to compute the k-anonymity of a
@@ -380,8 +385,10 @@ def callback(message):
     if not finished:
         print('No event received before the timeout. Please verify that the '
               'subscription provided is subscribed to the topic provided.')
+# [END dlp_k_anonymity]
 
 
+# [START dlp_l_diversity]
 def l_diversity_analysis(project, table_project_id, dataset_id, table_id,
                          topic_id, subscription_id, sensitive_attribute,
                          quasi_ids, timeout=300):
@@ -515,8 +522,10 @@ def callback(message):
     if not finished:
         print('No event received before the timeout. Please verify that the '
               'subscription provided is subscribed to the topic provided.')
+# [END dlp_l_diversity]
 
 
+# [START dlp_k_map]
 def k_map_estimate_analysis(project, table_project_id, dataset_id, table_id,
                             topic_id, subscription_id, quasi_ids, info_types,
                             region_code='US', timeout=300):
@@ -655,6 +664,7 @@ def callback(message):
     if not finished:
         print('No event received before the timeout. Please verify that the '
               'subscription provided is subscribed to the topic provided.')
+# [END dlp_k_map]
 
 
 if __name__ == '__main__':

dlp/risk_test.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from gcp_devrel.testing.flaky import flaky
 import google.cloud.pubsub
 
 import pytest
@@ -62,6 +63,7 @@ def subscription_id(topic_id):
     subscriber.delete_subscription(subscription_path)
 
 
+@flaky
 def test_numerical_risk_analysis(topic_id, subscription_id, capsys):
     risk.numerical_risk_analysis(
         GCLOUD_PROJECT,
@@ -76,6 +78,7 @@ def test_numerical_risk_analysis(topic_id, subscription_id, capsys):
     assert 'Value Range:' in out
 
 
+@flaky
 def test_categorical_risk_analysis_on_string_field(
         topic_id, subscription_id, capsys):
     risk.categorical_risk_analysis(
@@ -91,6 +94,7 @@ def test_categorical_risk_analysis_on_string_field(
     assert 'Most common value occurs' in out
 
 
+@flaky
 def test_categorical_risk_analysis_on_number_field(
         topic_id, subscription_id, capsys):
     risk.categorical_risk_analysis(
@@ -106,6 +110,7 @@ def test_categorical_risk_analysis_on_number_field(
     assert 'Most common value occurs' in out
 
 
+@flaky
 def test_k_anonymity_analysis_single_field(topic_id, subscription_id, capsys):
     risk.k_anonymity_analysis(
         GCLOUD_PROJECT,
@@ -121,6 +126,7 @@ def test_k_anonymity_analysis_single_field(topic_id, subscription_id, capsys):
     assert 'Class size:' in out
 
 
+@flaky
 def test_k_anonymity_analysis_multiple_fields(topic_id, subscription_id,
                                               capsys):
     risk.k_anonymity_analysis(
@@ -137,6 +143,7 @@ def test_k_anonymity_analysis_multiple_fields(topic_id, subscription_id,
     assert 'Class size:' in out
 
 
+@flaky
 def test_l_diversity_analysis_single_field(topic_id, subscription_id, capsys):
     risk.l_diversity_analysis(
         GCLOUD_PROJECT,
@@ -154,6 +161,7 @@ def test_l_diversity_analysis_single_field(topic_id, subscription_id, capsys):
     assert 'Sensitive value' in out
 
 
+@flaky
 def test_l_diversity_analysis_multiple_field(
         topic_id, subscription_id, capsys):
     risk.l_diversity_analysis(
@@ -172,6 +180,7 @@ def test_l_diversity_analysis_multiple_field(
     assert 'Sensitive value' in out
 
 
+@flaky
 def test_k_map_estimate_analysis_single_field(
         topic_id, subscription_id, capsys):
     risk.k_map_estimate_analysis(
@@ -190,6 +199,7 @@ def test_k_map_estimate_analysis_single_field(
     assert 'Values' in out
 
 
+@flaky
 def test_k_map_estimate_analysis_multiple_field(
         topic_id, subscription_id, capsys):
     risk.k_map_estimate_analysis(
@@ -208,6 +218,7 @@ def test_k_map_estimate_analysis_multiple_field(
     assert 'Values' in out
 
 
+@flaky
 def test_k_map_estimate_analysis_quasi_ids_info_types_equal(
         topic_id, subscription_id):
     with pytest.raises(ValueError):