From 9ccab779d5da04db9cd6811008b26598295dd52d Mon Sep 17 00:00:00 2001
From: Bu Sun Kim <8822365+busunkim96@users.noreply.github.com>
Date: Tue, 5 Oct 2021 04:20:10 -0600
Subject: [PATCH] samples: increase timeout, catch
 concurrent.futures.TimeoutError (#266)

* samples: catch concurrent.futures.TimeoutError

* fix: use default timeout in samples file
---
 .../google-cloud-dlp/samples/snippets/risk.py     | 15 ++++++++++-----
 .../samples/snippets/risk_test.py                 | 12 ------------
 2 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/packages/google-cloud-dlp/samples/snippets/risk.py b/packages/google-cloud-dlp/samples/snippets/risk.py
index 59f7362a57a0..a630eee260d8 100644
--- a/packages/google-cloud-dlp/samples/snippets/risk.py
+++ b/packages/google-cloud-dlp/samples/snippets/risk.py
@@ -48,6 +48,7 @@ def numerical_risk_analysis(
     Returns:
         None; the response from the API is printed to the terminal.
     """
+    import concurrent.futures
 
     # Import the client library.
     import google.cloud.dlp
@@ -116,7 +117,7 @@ def callback(message):
 
     try:
         subscription.result(timeout=timeout)
-    except TimeoutError:
+    except concurrent.futures.TimeoutError:
         print(
             "No event received before the timeout. Please verify that the "
             "subscription provided is subscribed to the topic provided."
@@ -156,6 +157,7 @@ def categorical_risk_analysis(
     Returns:
         None; the response from the API is printed to the terminal.
     """
+    import concurrent.futures
 
     # Import the client library.
     import google.cloud.dlp
@@ -237,7 +239,7 @@ def callback(message):
 
     try:
         subscription.result(timeout=timeout)
-    except TimeoutError:
+    except concurrent.futures.TimeoutError:
         print(
             "No event received before the timeout. Please verify that the "
             "subscription provided is subscribed to the topic provided."
@@ -277,6 +279,7 @@ def k_anonymity_analysis(
     Returns:
         None; the response from the API is printed to the terminal.
     """
+    import concurrent.futures
 
     # Import the client library.
     import google.cloud.dlp
@@ -367,7 +370,7 @@ def callback(message):
 
     try:
         subscription.result(timeout=timeout)
-    except TimeoutError:
+    except concurrent.futures.TimeoutError:
         print(
             "No event received before the timeout. Please verify that the "
             "subscription provided is subscribed to the topic provided."
@@ -409,6 +412,7 @@ def l_diversity_analysis(
     Returns:
         None; the response from the API is printed to the terminal.
     """
+    import concurrent.futures
 
     # Import the client library.
     import google.cloud.dlp
@@ -509,7 +513,7 @@ def callback(message):
 
     try:
         subscription.result(timeout=timeout)
-    except TimeoutError:
+    except concurrent.futures.TimeoutError:
         print(
             "No event received before the timeout. Please verify that the "
             "subscription provided is subscribed to the topic provided."
@@ -558,6 +562,7 @@ def k_map_estimate_analysis(
     Returns:
         None; the response from the API is printed to the terminal.
     """
+    import concurrent.futures
 
     # Import the client library.
     import google.cloud.dlp
@@ -659,7 +664,7 @@ def callback(message):
 
     try:
         subscription.result(timeout=timeout)
-    except TimeoutError:
+    except concurrent.futures.TimeoutError:
         print(
             "No event received before the timeout. Please verify that the "
             "subscription provided is subscribed to the topic provided."
diff --git a/packages/google-cloud-dlp/samples/snippets/risk_test.py b/packages/google-cloud-dlp/samples/snippets/risk_test.py
index ef2ac3270019..a57abda35192 100644
--- a/packages/google-cloud-dlp/samples/snippets/risk_test.py
+++ b/packages/google-cloud-dlp/samples/snippets/risk_test.py
@@ -36,8 +36,6 @@
 BIGQUERY_TABLE_ID = "dlp_test_table" + UNIQUE_STRING
 BIGQUERY_HARMFUL_TABLE_ID = "harmful" + UNIQUE_STRING
 
-TIMEOUT = 120  # 2 minutes
-
 
 # Create new custom topic/subscription
 # We observe sometimes all the tests in this file fail. In a
@@ -173,7 +171,6 @@ def test_numerical_risk_analysis(topic_id, subscription_id, bigquery_project, ca
         NUMERIC_FIELD,
         topic_id,
         subscription_id,
-        timeout=TIMEOUT,
     )
 
     out, _ = capsys.readouterr()
@@ -192,7 +189,6 @@ def test_categorical_risk_analysis_on_string_field(
         UNIQUE_FIELD,
         topic_id,
         subscription_id,
-        timeout=TIMEOUT,
     )
 
     out, _ = capsys.readouterr()
@@ -211,7 +207,6 @@ def test_categorical_risk_analysis_on_number_field(
         NUMERIC_FIELD,
         topic_id,
         subscription_id,
-        timeout=TIMEOUT,
     )
 
     out, _ = capsys.readouterr()
@@ -230,7 +225,6 @@ def test_k_anonymity_analysis_single_field(
         topic_id,
         subscription_id,
         [NUMERIC_FIELD],
-        timeout=TIMEOUT,
     )
 
     out, _ = capsys.readouterr()
@@ -250,7 +244,6 @@ def test_k_anonymity_analysis_multiple_fields(
         topic_id,
         subscription_id,
         [NUMERIC_FIELD, REPEATED_FIELD],
-        timeout=TIMEOUT,
     )
 
     out, _ = capsys.readouterr()
@@ -271,7 +264,6 @@ def test_l_diversity_analysis_single_field(
         subscription_id,
         UNIQUE_FIELD,
         [NUMERIC_FIELD],
-        timeout=TIMEOUT,
     )
 
     out, _ = capsys.readouterr()
@@ -293,7 +285,6 @@ def test_l_diversity_analysis_multiple_field(
         subscription_id,
         UNIQUE_FIELD,
         [NUMERIC_FIELD, REPEATED_FIELD],
-        timeout=TIMEOUT,
     )
 
     out, _ = capsys.readouterr()
@@ -315,7 +306,6 @@ def test_k_map_estimate_analysis_single_field(
         subscription_id,
         [NUMERIC_FIELD],
         ["AGE"],
-        timeout=TIMEOUT,
     )
 
     out, _ = capsys.readouterr()
@@ -337,7 +327,6 @@ def test_k_map_estimate_analysis_multiple_field(
         subscription_id,
         [NUMERIC_FIELD, STRING_BOOLEAN_FIELD],
         ["AGE", "GENDER"],
-        timeout=TIMEOUT,
     )
 
     out, _ = capsys.readouterr()
@@ -360,5 +349,4 @@ def test_k_map_estimate_analysis_quasi_ids_info_types_equal(
             subscription_id,
             [NUMERIC_FIELD, STRING_BOOLEAN_FIELD],
             ["AGE"],
-            timeout=TIMEOUT,
         )