fix: making test more resistant to concurent running tests (GoogleCloudPlatform#11842)

Stepan Rasputny · web-flow · commit 15f85630a8fa · 2024-06-17T14:39:49.000+02:00
diff --git a/iam/cloud-client/snippets/test_roles.py b/iam/cloud-client/snippets/test_roles.py
@@ -14,6 +14,8 @@
 
 import re
 
+import backoff
+from google.api_core.exceptions import Aborted, InvalidArgument
 import google.auth
 from google.cloud.iam_admin_v1 import GetRoleRequest, IAMClient, ListRolesRequest, Role
 import pytest
@@ -67,6 +69,7 @@ def test_list_roles(capsys: "pytest.CaptureFixture[str]", iam_role: str) -> None
     assert re.search(iam_role, out)
 
 
+@backoff.on_exception(backoff.expo, Aborted, max_tries=3)
 def test_edit_role(iam_role: str) -> None:
     client = IAMClient()
     name = f"projects/{PROJECT}/roles/{iam_role}"
@@ -79,6 +82,7 @@ def test_edit_role(iam_role: str) -> None:
     assert updated_role.title == title
 
 
+@backoff.on_exception(backoff.expo, InvalidArgument, max_tries=3)
 def test_disable_role(capsys: "pytest.CaptureFixture[str]", iam_role: str) -> None:
     disable_role(PROJECT, iam_role)
     client = IAMClient()
diff --git a/iam/cloud-client/snippets/test_service_account.py b/iam/cloud-client/snippets/test_service_account.py
@@ -16,7 +16,7 @@
 import uuid
 
 import backoff
-from google.api_core.exceptions import InvalidArgument
+from google.api_core.exceptions import InvalidArgument, NotFound
 import google.auth
 from google.iam.v1 import policy_pb2
 import pytest
@@ -92,7 +92,7 @@ def test_service_account_set_policy(service_account: str) -> None:
 
     try:
         new_policy = set_service_account_iam_policy(PROJECT, service_account, policy)
-    except InvalidArgument:
+    except (InvalidArgument, NotFound):
         pytest.skip("Service account was removed from outside, skipping")
 
     binding_found = False
@@ -108,7 +108,7 @@ def test_service_account_rename(service_account: str) -> None:
     new_name = "New Name"
     try:
         account = rename_service_account(PROJECT, service_account, new_name)
-    except InvalidArgument:
+    except (InvalidArgument, NotFound):
         pytest.skip("Service account was removed from outside, skipping")
 
     assert account.display_name == new_name
