wrapping flaky parts into backoff

iam/cloud-client/snippets/get_policy.py

@@ -19,7 +19,7 @@
 from google.iam.v1 import iam_policy_pb2, policy_pb2
 
 
-def get_policy(project_id: str) -> policy_pb2.Policy:
+def get_project_policy(project_id: str) -> policy_pb2.Policy:
     """
     Get policy for project.
 
@@ -44,4 +44,4 @@ def get_policy(project_id: str) -> policy_pb2.Policy:
     # Your Google Cloud project ID.
     project_id = "test-project-id"
 
-    policy = get_policy(project_id)
+    policy = get_project_policy(project_id)

iam/cloud-client/snippets/modify_policy_add_member.py

@@ -14,8 +14,8 @@
 
 # [START iam_modify_policy_add_member]
 from google.iam.v1 import policy_pb2
-from snippets.get_policy import get_policy
-from snippets.set_policy import set_policy
+from snippets.get_policy import get_project_policy
+from snippets.set_policy import set_project_policy
 
 
 def modify_policy_add_member(
@@ -37,14 +37,14 @@ def modify_policy_add_member(
         * deleted:group:{emailid}?uid={uniqueid}
         * domain:{domain}
     """
-    policy = get_policy(project_id)
+    policy = get_project_policy(project_id)
 
     for bind in policy.bindings:
         if bind.role == role:
             bind.members.append(member)
             break
 
-    return set_policy(project_id, policy)
+    return set_project_policy(project_id, policy)
 
 
 # [END iam_modify_policy_add_member]

iam/cloud-client/snippets/modify_policy_remove_member.py

@@ -14,8 +14,8 @@
 
 # [START iam_modify_policy_remove_member]
 from google.iam.v1 import policy_pb2
-from snippets.get_policy import get_policy
-from snippets.set_policy import set_policy
+from snippets.get_policy import get_project_policy
+from snippets.set_policy import set_project_policy
 
 
 def modify_policy_remove_member(
@@ -37,15 +37,15 @@ def modify_policy_remove_member(
         * deleted:group:{emailid}?uid={uniqueid}
         * domain:{domain}
     """
-    policy = get_policy(project_id)
+    policy = get_project_policy(project_id)
 
     for bind in policy.bindings:
         if bind.role == role:
             if member in bind.members:
                 bind.members.remove(member)
             break
 
-    return set_policy(project_id, policy)
+    return set_project_policy(project_id, policy)
 
 
 # [END iam_modify_policy_remove_member]

iam/cloud-client/snippets/set_policy.py

@@ -19,7 +19,7 @@
 from google.iam.v1 import iam_policy_pb2, policy_pb2
 
 
-def set_policy(project_id: str, policy: policy_pb2.Policy) -> policy_pb2.Policy:
+def set_project_policy(project_id: str, policy: policy_pb2.Policy) -> policy_pb2.Policy:
     """
     Set policy for project. Pay attention that previous state will be completely rewritten.
     If you want to update only part of the policy follow the approach read->modify->write.
@@ -28,15 +28,22 @@ def set_policy(project_id: str, policy: policy_pb2.Policy) -> policy_pb2.Policy:
     project_id: ID or number of the Google Cloud project you want to use.
     policy: Policy which has to be set.
     """
-
     client = resourcemanager_v3.ProjectsClient()
+
+    request = iam_policy_pb2.GetIamPolicyRequest()
+    request.resource = f"projects/{project_id}"
+    current_policy = client.get_iam_policy(request)
+
+    # Etag should as fresh as possible to lower chance of collisions
+    policy.ClearField("etag")
+    current_policy.CopyFrom(policy)
     request = iam_policy_pb2.SetIamPolicyRequest()
     request.resource = f"projects/{project_id}"
 
     # request.etag field also will be merged which means you are secured from collision,
-    # but it means that request may fail and you need to leverage exponential reties approach
+    # but it means that request may fail and you need to leverage exponential retries approach
     # to be sure policy has been updated.
-    request.policy.MergeFrom(policy)
+    request.policy.CopyFrom(current_policy)
 
     policy = client.set_iam_policy(request)
     return policy
@@ -52,14 +59,10 @@ def set_policy(project_id: str, policy: policy_pb2.Policy) -> policy_pb2.Policy:
     # Your Google Cloud project ID.
     project_id = "test-project-id"
 
-    bindings = [
-        {
-            "role": "roles/viewer",
-            "members": [
-                "serviceAccount:test-service-account@test-project-id.iam.gserviceaccount.com",
-            ],
-            "condition": {},
-        },
-    ]
-
-    set_policy(project_id, bindings)
+    new_policy = policy_pb2.Policy()
+    binding = policy_pb2.Binding()
+    binding.role = "roles/viewer"
+    binding.members.append(f"serviceAccount:test-service-account@{project_id}.iam.gserviceaccount.com")
+    new_policy.bindings.append(binding)
+
+    set_project_policy(project_id, new_policy)

iam/cloud-client/snippets/test_project_policies.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import re
+from typing import Callable, Union
 import uuid
 
 import backoff
@@ -22,11 +23,11 @@
 import pytest
 from snippets.create_service_account import create_service_account
 from snippets.delete_service_account import delete_service_account
-from snippets.get_policy import get_policy
+from snippets.get_policy import get_project_policy
 from snippets.modify_policy_add_member import modify_policy_add_member
 from snippets.modify_policy_remove_member import modify_policy_remove_member
 from snippets.query_testable_permissions import query_testable_permissions
-from snippets.set_policy import set_policy
+from snippets.set_policy import set_project_policy
 
 PROJECT = google.auth.default()[1]
 
@@ -50,24 +51,23 @@ def service_account(capsys: "pytest.CaptureFixture[str]") -> str:
 @pytest.fixture
 def project_policy() -> policy_pb2.Policy:
     try:
-        policy = get_policy(PROJECT)
+        policy = get_project_policy(PROJECT)
         policy_copy = policy_pb2.Policy()
         policy_copy.CopyFrom(policy)
         yield policy_copy
     finally:
-        policy.ClearField("etag")
-        updated_policy = set_policy(PROJECT, policy)
+        updated_policy = execute_wrapped(set_project_policy, PROJECT, policy)
 
         updated_policy.ClearField("etag")
         assert updated_policy == policy
 
 
 @backoff.on_exception(backoff.expo, Aborted, max_tries=3)
-def policy_setter(project: str, policy: policy_pb2.Policy) -> policy_pb2.Policy:
-    return set_policy(project, policy)
+def execute_wrapped(func: Callable, *args: Union[str, policy_pb2.Policy]) -> policy_pb2.Policy:
+    return func(*args)
 
 
-def test_set_policy(project_policy: policy_pb2.Policy) -> None:
+def test_set_project_policy(project_policy: policy_pb2.Policy) -> None:
     role = "roles/viewer"
     test_binding = policy_pb2.Binding()
     test_binding.role = role
@@ -78,7 +78,7 @@ def test_set_policy(project_policy: policy_pb2.Policy) -> None:
     )
     project_policy.bindings.append(test_binding)
 
-    policy = policy_setter(PROJECT, project_policy)
+    policy = execute_wrapped(set_project_policy, PROJECT, project_policy)
 
     binding_found = False
     for bind in policy.bindings:
@@ -101,7 +101,7 @@ def test_modify_policy_add_member(
     )
     project_policy.bindings.append(test_binding)
 
-    policy = policy_setter(PROJECT, project_policy)
+    policy = execute_wrapped(set_project_policy, PROJECT, project_policy)
     binding_found = False
     for bind in policy.bindings:
         if bind.role == test_binding.role:
@@ -110,7 +110,7 @@ def test_modify_policy_add_member(
     assert binding_found
 
     member = f"serviceAccount:{service_account}"
-    policy = modify_policy_add_member(PROJECT, role, member)
+    policy = execute_wrapped(modify_policy_add_member, PROJECT, role, member)
 
     member_added = False
     for bind in policy.bindings:
@@ -135,7 +135,7 @@ def test_modify_policy_remove_member(
     )
     project_policy.bindings.append(test_binding)
 
-    policy = policy_setter(PROJECT, project_policy)
+    policy = execute_wrapped(set_project_policy, PROJECT, project_policy)
 
     binding_found = False
     for bind in policy.bindings:
@@ -144,7 +144,7 @@ def test_modify_policy_remove_member(
             break
     assert binding_found
 
-    policy = modify_policy_remove_member(PROJECT, role, member)
+    policy = execute_wrapped(modify_policy_remove_member, PROJECT, role, member)
 
     member_removed = False
     for bind in policy.bindings: