Cleanup iam tests (GoogleCloudPlatform#2701)

* clean up access & add optl policy version param

* Cleanup custom roles test

* Blacken and lint

Author: leahecole

iam/api-client/access.py

@@ -26,54 +26,67 @@
 
 
 # [START iam_get_policy]
-def get_policy(project_id):
+def get_policy(project_id, version=1):
     """Gets IAM policy for a project."""
 
     credentials = service_account.Credentials.from_service_account_file(
-        filename=os.environ['GOOGLE_APPLICATION_CREDENTIALS'],
-        scopes=['https://www.googleapis.com/auth/cloud-platform'])
+        filename=os.environ["GOOGLE_APPLICATION_CREDENTIALS"],
+        scopes=["https://www.googleapis.com/auth/cloud-platform"],
+    )
     service = googleapiclient.discovery.build(
-        'cloudresourcemanager', 'v1', credentials=credentials)
-    policy = service.projects().getIamPolicy(
-        resource=project_id, body={}).execute()
+        "cloudresourcemanager", "v1", credentials=credentials
+    )
+    policy = (
+        service.projects()
+        .getIamPolicy(
+            resource=project_id,
+            body={"options": {"requestedPolicyVersion": version}},
+        )
+        .execute()
+    )
     print(policy)
     return policy
+
+
 # [END iam_get_policy]
 
 
 # [START iam_modify_policy_add_member]
 def modify_policy_add_member(policy, role, member):
     """Adds a new member to a role binding."""
 
-    binding = next(b for b in policy['bindings'] if b['role'] == role)
-    binding['members'].append(member)
+    binding = next(b for b in policy["bindings"] if b["role"] == role)
+    binding["members"].append(member)
     print(binding)
     return policy
+
+
 # [END iam_modify_policy_add_member]
 
 
 # [START iam_modify_policy_add_role]
 def modify_policy_add_role(policy, role, member):
     """Adds a new role binding to a policy."""
 
-    binding = {
-        'role': role,
-        'members': [member]
-    }
-    policy['bindings'].append(binding)
+    binding = {"role": role, "members": [member]}
+    policy["bindings"].append(binding)
     print(policy)
     return policy
+
+
 # [END iam_modify_policy_add_role]
 
 
 # [START iam_modify_policy_remove_member]
 def modify_policy_remove_member(policy, role, member):
     """Removes a  member from a role binding."""
-    binding = next(b for b in policy['bindings'] if b['role'] == role)
-    if 'members' in binding and member in binding['members']:
-        binding['members'].remove(member)
+    binding = next(b for b in policy["bindings"] if b["role"] == role)
+    if "members" in binding and member in binding["members"]:
+        binding["members"].remove(member)
     print(binding)
     return policy
+
+
 # [END iam_modify_policy_remove_member]
 
 
@@ -82,17 +95,22 @@ def set_policy(project_id, policy):
     """Sets IAM policy for a project."""
 
     credentials = service_account.Credentials.from_service_account_file(
-        filename=os.environ['GOOGLE_APPLICATION_CREDENTIALS'],
-        scopes=['https://www.googleapis.com/auth/cloud-platform'])
+        filename=os.environ["GOOGLE_APPLICATION_CREDENTIALS"],
+        scopes=["https://www.googleapis.com/auth/cloud-platform"],
+    )
     service = googleapiclient.discovery.build(
-        'cloudresourcemanager', 'v1', credentials=credentials)
-
-    policy = service.projects().setIamPolicy(
-        resource=project_id, body={
-            'policy': policy
-        }).execute()
+        "cloudresourcemanager", "v1", credentials=credentials
+    )
+
+    policy = (
+        service.projects()
+        .setIamPolicy(resource=project_id, body={"policy": policy})
+        .execute()
+    )
     print(policy)
     return policy
+
+
 # [END iam_set_policy]
 
 
@@ -101,86 +119,94 @@ def test_permissions(project_id):
     """Tests IAM permissions of the caller"""
 
     credentials = service_account.Credentials.from_service_account_file(
-        filename=os.environ['GOOGLE_APPLICATION_CREDENTIALS'],
-        scopes=['https://www.googleapis.com/auth/cloud-platform'])
+        filename=os.environ["GOOGLE_APPLICATION_CREDENTIALS"],
+        scopes=["https://www.googleapis.com/auth/cloud-platform"],
+    )
     service = googleapiclient.discovery.build(
-        'cloudresourcemanager', 'v1', credentials=credentials)
+        "cloudresourcemanager", "v1", credentials=credentials
+    )
 
     permissions = {
         "permissions": [
             "resourcemanager.projects.get",
-            "resourcemanager.projects.delete"
+            "resourcemanager.projects.delete",
         ]
     }
 
     request = service.projects().testIamPermissions(
-        resource=project_id, body=permissions)
+        resource=project_id, body=permissions
+    )
     returnedPermissions = request.execute()
     print(returnedPermissions)
     return returnedPermissions
+
+
 # [END iam_test_permissions]
 
 
 def main():
     parser = argparse.ArgumentParser(
         description=__doc__,
-        formatter_class=argparse.RawDescriptionHelpFormatter)
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
 
-    subparsers = parser.add_subparsers(dest='command')
+    subparsers = parser.add_subparsers(dest="command")
 
     # Get
-    get_parser = subparsers.add_parser(
-        'get', help=get_policy.__doc__)
-    get_parser.add_argument('project_id')
+    get_parser = subparsers.add_parser("get", help=get_policy.__doc__)
+    get_parser.add_argument("project_id")
 
     # Modify: add member
     modify_member_parser = subparsers.add_parser(
-        'modify_member', help=get_policy.__doc__)
-    modify_member_parser.add_argument('project_id')
-    modify_member_parser.add_argument('role')
-    modify_member_parser.add_argument('member')
+        "modify_member", help=get_policy.__doc__
+    )
+    modify_member_parser.add_argument("project_id")
+    modify_member_parser.add_argument("role")
+    modify_member_parser.add_argument("member")
 
     # Modify: add role
     modify_role_parser = subparsers.add_parser(
-        'modify_role', help=get_policy.__doc__)
-    modify_role_parser.add_argument('project_id')
-    modify_role_parser.add_argument('project_id')
-    modify_role_parser.add_argument('role')
-    modify_role_parser.add_argument('member')
+        "modify_role", help=get_policy.__doc__
+    )
+    modify_role_parser.add_argument("project_id")
+    modify_role_parser.add_argument("project_id")
+    modify_role_parser.add_argument("role")
+    modify_role_parser.add_argument("member")
 
     # Modify: remove member
     modify_member_parser = subparsers.add_parser(
-        'modify_member', help=get_policy.__doc__)
-    modify_member_parser.add_argument('project_id')
-    modify_member_parser.add_argument('role')
-    modify_member_parser.add_argument('member')
+        "modify_member", help=get_policy.__doc__
+    )
+    modify_member_parser.add_argument("project_id")
+    modify_member_parser.add_argument("role")
+    modify_member_parser.add_argument("member")
 
     # Set
-    set_parser = subparsers.add_parser(
-        'set', help=set_policy.__doc__)
-    set_parser.add_argument('project_id')
-    set_parser.add_argument('policy')
+    set_parser = subparsers.add_parser("set", help=set_policy.__doc__)
+    set_parser.add_argument("project_id")
+    set_parser.add_argument("policy")
 
     # Test permissions
     test_permissions_parser = subparsers.add_parser(
-        'test_permissions', help=get_policy.__doc__)
-    test_permissions_parser.add_argument('project_id')
+        "test_permissions", help=get_policy.__doc__
+    )
+    test_permissions_parser.add_argument("project_id")
 
     args = parser.parse_args()
 
-    if args.command == 'get':
+    if args.command == "get":
         get_policy(args.project_id)
-    elif args.command == 'set':
+    elif args.command == "set":
         set_policy(args.project_id, args.policy)
-    elif args.command == 'add_member':
+    elif args.command == "add_member":
         modify_policy_add_member(args.policy, args.role, args.member)
-    elif args.command == 'remove_member':
+    elif args.command == "remove_member":
         modify_policy_remove_member(args.policy, args.role, args.member)
-    elif args.command == 'add_binding':
+    elif args.command == "add_binding":
         modify_policy_add_role(args.policy, args.role, args.member)
-    elif args.command == 'test_permissions':
+    elif args.command == "test_permissions":
         test_permissions(args.project_id)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

iam/api-client/access_test.py

@@ -13,47 +13,64 @@
 # limitations under the License.
 
 import os
+import pytest
 import random
 
 import access
 import service_accounts
 
+# Setting up variables for testing
+GCLOUD_PROJECT = os.environ["GCLOUD_PROJECT"]
 
-def test_access(capsys):
-    # Setting up variables for testing
-    project_id = os.environ['GCLOUD_PROJECT']
+# specifying a sample role to be assigned
+GCP_ROLE = "roles/owner"
 
-    # specifying a sample role to be assigned
-    gcp_role = 'roles/owner'
 
+@pytest.fixture(scope="module")
+def test_member():
     # section to create service account to test policy updates.
     rand = str(random.randint(0, 1000))
-    name = 'python-test-' + rand
-    email = name + '@' + project_id + '.iam.gserviceaccount.com'
-    member = 'serviceAccount:' + email
+    name = "python-test-" + rand
+    email = name + "@" + GCLOUD_PROJECT + ".iam.gserviceaccount.com"
+    member = "serviceAccount:" + email
     service_accounts.create_service_account(
-        project_id, name, 'Py Test Account')
+        GCLOUD_PROJECT, name, "Py Test Account"
+    )
 
-    policy = access.get_policy(project_id)
-    out, _ = capsys.readouterr()
-    assert u'etag' in out
+    yield member
+
+    # deleting the service account created above
+    service_accounts.delete_service_account(email)
 
-    policy = access.modify_policy_add_role(policy, gcp_role, member)
+
+def test_get_policy(capsys):
+    access.get_policy(GCLOUD_PROJECT, version=3)
     out, _ = capsys.readouterr()
-    assert u'etag' in out
+    assert u"etag" in out
+
 
-    policy = access.modify_policy_remove_member(policy, gcp_role, member)
+def test_modify_policy_add_role(test_member, capsys):
+    policy = access.get_policy(GCLOUD_PROJECT, version=3)
+    access.modify_policy_add_role(policy, GCLOUD_PROJECT, test_member)
     out, _ = capsys.readouterr()
-    assert 'iam.gserviceaccount.com' in out
+    assert u"etag" in out
 
-    policy = access.set_policy(project_id, policy)
+
+def test_modify_policy_remove_member(test_member, capsys):
+    policy = access.get_policy(GCLOUD_PROJECT, version=3)
+    access.modify_policy_remove_member(policy, GCP_ROLE, test_member)
     out, _ = capsys.readouterr()
-    assert u'etag' in out
+    assert "iam.gserviceaccount.com" in out
+
 
-    access.test_permissions(project_id)
+def test_set_policy(capsys):
+    policy = access.get_policy(GCLOUD_PROJECT, version=3)
+    access.set_policy(GCLOUD_PROJECT, policy)
     out, _ = capsys.readouterr()
-    assert u'permissions' in out
+    assert u"etag" in out
 
-    # deleting the service account created above
-    service_accounts.delete_service_account(
-        email)
+
+def test_permissions(capsys):
+    access.test_permissions(GCLOUD_PROJECT)
+    out, _ = capsys.readouterr()
+    assert u"permissions" in out