chore(scc): use backoff in tests (#10444)

* chore(scc): use backoff in tests * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
GoogleCloudPlatform · Jan 17, 2024 · 24a6f7d · 24a6f7d
1 parent 5820cd4
commit 24a6f7d
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 6 deletions.
diff --git a/securitycenter/snippets/snippets_bigquery_export_test.py b/securitycenter/snippets/snippets_bigquery_export_test.py
@@ -22,7 +22,7 @@
 
 from _pytest.capture import CaptureFixture
 import backoff
-from google.api_core.exceptions import ServiceUnavailable
+from google.api_core.exceptions import InternalServerError, NotFound, ServiceUnavailable
 import pytest
 
 import snippets_bigquery_export
@@ -50,7 +50,9 @@ def bigquery_export_id():
     delete_bigquery_dataset(BIGQUERY_DATASET_ID)
 
 
-@backoff.on_exception(backoff.expo, ServiceUnavailable, max_tries=3)
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def create_bigquery_dataset(dataset_id: str):
     from google.cloud import bigquery
 
@@ -63,7 +65,9 @@ def create_bigquery_dataset(dataset_id: str):
     print(f"Dataset {dataset.dataset_id} created.")
 
 
-@backoff.on_exception(backoff.expo, ServiceUnavailable, max_tries=3)
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def delete_bigquery_dataset(dataset_id: str):
     from google.cloud import bigquery
 
@@ -72,7 +76,9 @@ def delete_bigquery_dataset(dataset_id: str):
     print(f"Dataset {dataset_id} deleted.")
 
 
-@backoff.on_exception(backoff.expo, ServiceUnavailable, max_tries=3)
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_get_bigquery_export(capsys: CaptureFixture, bigquery_export_id: str):
     snippets_bigquery_export.get_bigquery_export(
         f"projects/{PROJECT_ID}", bigquery_export_id
@@ -85,15 +91,19 @@ def test_get_bigquery_export(capsys: CaptureFixture, bigquery_export_id: str):
     assert re.search(f"bigQueryExports/{bigquery_export_id}", out)
 
 
-@backoff.on_exception(backoff.expo, ServiceUnavailable, max_tries=3)
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_list_bigquery_exports(capsys: CaptureFixture, bigquery_export_id: str):
     snippets_bigquery_export.list_bigquery_exports(f"projects/{PROJECT_ID}")
     out, _ = capsys.readouterr()
     assert re.search("Listing BigQuery exports:", out)
     assert re.search(bigquery_export_id, out)
 
 
-@backoff.on_exception(backoff.expo, ServiceUnavailable, max_tries=3)
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_update_bigquery_exports(capsys: CaptureFixture, bigquery_export_id: str):
     export_filter = 'severity="MEDIUM"'
     snippets_bigquery_export.update_bigquery_export(

diff --git a/securitycenter/snippets/snippets_mute_config_test.py b/securitycenter/snippets/snippets_mute_config_test.py
@@ -18,6 +18,8 @@
 import uuid
 
 from _pytest.capture import CaptureFixture
+import backoff
+from google.api_core.exceptions import InternalServerError, NotFound, ServiceUnavailable
 from google.cloud import securitycenter
 from google.cloud.securitycenter_v1.services.security_center.pagers import (
     ListFindingsPager,
@@ -74,6 +76,9 @@ def list_all_findings(source_name) -> ListFindingsPager:
     return client.list_findings(request={"parent": source_name})
 
 
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_get_mute_rule(capsys: CaptureFixture, mute_rule):
     snippets_mute_config.get_mute_rule(
         f"projects/{PROJECT_ID}/muteConfigs/{mute_rule.get('create')}"
@@ -83,13 +88,19 @@ def test_get_mute_rule(capsys: CaptureFixture, mute_rule):
     assert re.search(mute_rule.get("create"), out)
 
 
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_list_mute_rules(capsys: CaptureFixture, mute_rule):
     snippets_mute_config.list_mute_rules(f"projects/{PROJECT_ID}")
     out, _ = capsys.readouterr()
     assert re.search(mute_rule.get("create"), out)
     assert re.search(mute_rule.get("update"), out)
 
 
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_update_mute_rule(capsys: CaptureFixture, mute_rule):
     snippets_mute_config.update_mute_rule(
         f"projects/{PROJECT_ID}/muteConfigs/{mute_rule.get('update')}"
@@ -101,20 +112,29 @@ def test_update_mute_rule(capsys: CaptureFixture, mute_rule):
     assert re.search("Updated mute config description", out)
 
 
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_set_mute_finding(capsys: CaptureFixture, finding):
     finding_path = finding.get("finding1")
     snippets_mute_config.set_mute_finding(finding_path)
     out, _ = capsys.readouterr()
     assert re.search("Mute value for the finding: MUTED", out)
 
 
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_set_unmute_finding(capsys: CaptureFixture, finding):
     finding_path = finding.get("finding1")
     snippets_mute_config.set_unmute_finding(finding_path)
     out, _ = capsys.readouterr()
     assert re.search("Mute value for the finding: UNMUTED", out)
 
 
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_bulk_mute_findings(capsys: CaptureFixture, finding):
     # Mute findings that belong to this project.
     snippets_mute_config.bulk_mute_findings(

diff --git a/securitycenter/snippets/snippets_notification_test.py b/securitycenter/snippets/snippets_notification_test.py
@@ -18,6 +18,8 @@
 import os
 import uuid
 
+import backoff
+from google.api_core.exceptions import InternalServerError, NotFound, ServiceUnavailable
 from google.cloud import securitycenter as securitycenter
 import pytest
 
@@ -108,6 +110,9 @@ def deleted_notification_config():
     return created_notification_config
 
 
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_create_notification_config():
     created_notification_config = (
         snippets_notification_configs.create_notification_config(
@@ -119,26 +124,38 @@ def test_create_notification_config():
     cleanup_notification_config(CREATE_CONFIG_ID)
 
 
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_delete_notification_config(deleted_notification_config):
     assert snippets_notification_configs.delete_notification_config(
         f"organizations/{ORG_ID}", DELETE_CONFIG_ID
     )
 
 
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_get_notification_config(new_notification_config_for_get):
     retrieved_config = snippets_notification_configs.get_notification_config(
         f"organizations/{ORG_ID}", GET_CONFIG_ID
     )
     assert retrieved_config is not None
 
 
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_list_notification_configs():
     iterator = snippets_notification_configs.list_notification_configs(
         f"organizations/{ORG_ID}"
     )
     assert iterator is not None
 
 
+@backoff.on_exception(
+    backoff.expo, (InternalServerError, ServiceUnavailable, NotFound), max_tries=3
+)
 def test_update_notification_config(new_notification_config_for_update):
     updated_config = snippets_notification_configs.update_notification_config(
         f"organizations/{ORG_ID}", UPDATE_CONFIG_ID, PUBSUB_TOPIC