samples: increase timeout, catch concurrent.futures.TimeoutError (#266)

* samples: catch concurrent.futures.TimeoutError * fix: use default timeout in samples file
GoogleCloudPlatform · Oct 5, 2021 · bb011e1 · bb011e1
1 parent 26aea2f
commit bb011e1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 17 deletions.
diff --git a/dlp/snippets/risk.py b/dlp/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/dlp/snippets/risk_test.py b/dlp/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,
         )