Skip to content

Commit

Permalink
Pub/Sub: update how to test with mock [(#2555)](GoogleCloudPlatform/p…
Browse files Browse the repository at this point in the history
…ython-docs-samples#2555)

* Update test with mock
* Clean up resources after tests
* Use unique resource names avoid test failures
* Delete subscriptions in cleanup phase
* Ensure unique topic name
* Update assert to remove bytestring notation
* Rewrite PubSubToGCS test using dataflow testing module
  • Loading branch information
anguillanneuf authored and busunkim96 committed Dec 11, 2019
1 parent aa682cc commit 6e46e5d
Show file tree
Hide file tree
Showing 12 changed files with 543 additions and 424 deletions.
115 changes: 62 additions & 53 deletions samples/snippets/iam.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,9 @@ def get_topic_policy(project, topic_name):

policy = client.get_iam_policy(topic_path)

print('Policy for topic {}:'.format(topic_path))
print("Policy for topic {}:".format(topic_path))
for binding in policy.bindings:
print('Role: {}, Members: {}'.format(binding.role, binding.members))
print("Role: {}, Members: {}".format(binding.role, binding.members))
# [END pubsub_get_topic_policy]


Expand All @@ -48,9 +48,9 @@ def get_subscription_policy(project, subscription_name):

policy = client.get_iam_policy(subscription_path)

print('Policy for subscription {}:'.format(subscription_path))
print("Policy for subscription {}:".format(subscription_path))
for binding in policy.bindings:
print('Role: {}, Members: {}'.format(binding.role, binding.members))
print("Role: {}, Members: {}".format(binding.role, binding.members))
# [END pubsub_get_subscription_policy]


Expand All @@ -63,20 +63,17 @@ def set_topic_policy(project, topic_name):
policy = client.get_iam_policy(topic_path)

# Add all users as viewers.
policy.bindings.add(
role='roles/pubsub.viewer',
members=['allUsers'])
policy.bindings.add(role="roles/pubsub.viewer", members=["allUsers"])

# Add a group as a publisher.
policy.bindings.add(
role='roles/pubsub.publisher',
members=['group:cloud-logs@google.com'])
role="roles/pubsub.publisher", members=["group:cloud-logs@google.com"]
)

# Set the policy
policy = client.set_iam_policy(topic_path, policy)

print('IAM policy for topic {} set: {}'.format(
topic_name, policy))
print("IAM policy for topic {} set: {}".format(topic_name, policy))
# [END pubsub_set_topic_policy]


Expand All @@ -89,20 +86,21 @@ def set_subscription_policy(project, subscription_name):
policy = client.get_iam_policy(subscription_path)

# Add all users as viewers.
policy.bindings.add(
role='roles/pubsub.viewer',
members=['allUsers'])
policy.bindings.add(role="roles/pubsub.viewer", members=["allUsers"])

# Add a group as an editor.
policy.bindings.add(
role='roles/editor',
members=['group:cloud-logs@google.com'])
role="roles/editor", members=["group:cloud-logs@google.com"]
)

# Set the policy
policy = client.set_iam_policy(subscription_path, policy)

print('IAM policy for subscription {} set: {}'.format(
subscription_name, policy))
print(
"IAM policy for subscription {} set: {}".format(
subscription_name, policy
)
)
# [END pubsub_set_subscription_policy]


Expand All @@ -112,16 +110,17 @@ def check_topic_permissions(project, topic_name):
client = pubsub_v1.PublisherClient()
topic_path = client.topic_path(project, topic_name)

permissions_to_check = [
'pubsub.topics.publish',
'pubsub.topics.update'
]
permissions_to_check = ["pubsub.topics.publish", "pubsub.topics.update"]

allowed_permissions = client.test_iam_permissions(
topic_path, permissions_to_check)
topic_path, permissions_to_check
)

print('Allowed permissions for topic {}: {}'.format(
topic_path, allowed_permissions))
print(
"Allowed permissions for topic {}: {}".format(
topic_path, allowed_permissions
)
)
# [END pubsub_test_topic_permissions]


Expand All @@ -132,63 +131,73 @@ def check_subscription_permissions(project, subscription_name):
subscription_path = client.subscription_path(project, subscription_name)

permissions_to_check = [
'pubsub.subscriptions.consume',
'pubsub.subscriptions.update'
"pubsub.subscriptions.consume",
"pubsub.subscriptions.update",
]

allowed_permissions = client.test_iam_permissions(
subscription_path, permissions_to_check)
subscription_path, permissions_to_check
)

print('Allowed permissions for subscription {}: {}'.format(
subscription_path, allowed_permissions))
print(
"Allowed permissions for subscription {}: {}".format(
subscription_path, allowed_permissions
)
)
# [END pubsub_test_subscription_permissions]


if __name__ == '__main__':
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument('project', help='Your Google Cloud project ID')
parser.add_argument("project", help="Your Google Cloud project ID")

subparsers = parser.add_subparsers(dest='command')
subparsers = parser.add_subparsers(dest="command")

get_topic_policy_parser = subparsers.add_parser(
'get-topic-policy', help=get_topic_policy.__doc__)
get_topic_policy_parser.add_argument('topic_name')
"get-topic-policy", help=get_topic_policy.__doc__
)
get_topic_policy_parser.add_argument("topic_name")

get_subscription_policy_parser = subparsers.add_parser(
'get-subscription-policy', help=get_subscription_policy.__doc__)
get_subscription_policy_parser.add_argument('subscription_name')
"get-subscription-policy", help=get_subscription_policy.__doc__
)
get_subscription_policy_parser.add_argument("subscription_name")

set_topic_policy_parser = subparsers.add_parser(
'set-topic-policy', help=set_topic_policy.__doc__)
set_topic_policy_parser.add_argument('topic_name')
"set-topic-policy", help=set_topic_policy.__doc__
)
set_topic_policy_parser.add_argument("topic_name")

set_subscription_policy_parser = subparsers.add_parser(
'set-subscription-policy', help=set_subscription_policy.__doc__)
set_subscription_policy_parser.add_argument('subscription_name')
"set-subscription-policy", help=set_subscription_policy.__doc__
)
set_subscription_policy_parser.add_argument("subscription_name")

check_topic_permissions_parser = subparsers.add_parser(
'check-topic-permissions', help=check_topic_permissions.__doc__)
check_topic_permissions_parser.add_argument('topic_name')
"check-topic-permissions", help=check_topic_permissions.__doc__
)
check_topic_permissions_parser.add_argument("topic_name")

check_subscription_permissions_parser = subparsers.add_parser(
'check-subscription-permissions',
help=check_subscription_permissions.__doc__)
check_subscription_permissions_parser.add_argument('subscription_name')
"check-subscription-permissions",
help=check_subscription_permissions.__doc__,
)
check_subscription_permissions_parser.add_argument("subscription_name")

args = parser.parse_args()

if args.command == 'get-topic-policy':
if args.command == "get-topic-policy":
get_topic_policy(args.project, args.topic_name)
elif args.command == 'get-subscription-policy':
elif args.command == "get-subscription-policy":
get_subscription_policy(args.project, args.subscription_name)
elif args.command == 'set-topic-policy':
elif args.command == "set-topic-policy":
set_topic_policy(args.project, args.topic_name)
elif args.command == 'set-subscription-policy':
elif args.command == "set-subscription-policy":
set_subscription_policy(args.project, args.subscription_name)
elif args.command == 'check-topic-permissions':
elif args.command == "check-topic-permissions":
check_topic_permissions(args.project, args.topic_name)
elif args.command == 'check-subscription-permissions':
elif args.command == "check-subscription-permissions":
check_subscription_permissions(args.project, args.subscription_name)
33 changes: 20 additions & 13 deletions samples/snippets/iam_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,23 +13,25 @@
# limitations under the License.

import os
import uuid

from google.cloud import pubsub_v1
import pytest

import iam

PROJECT = os.environ['GCLOUD_PROJECT']
TOPIC = 'iam-test-topic'
SUBSCRIPTION = 'iam-test-subscription'
UUID = uuid.uuid4().hex
PROJECT = os.environ["GCLOUD_PROJECT"]
TOPIC = "iam-test-topic-" + UUID
SUBSCRIPTION = "iam-test-subscription-" + UUID


@pytest.fixture(scope='module')
@pytest.fixture(scope="module")
def publisher_client():
yield pubsub_v1.PublisherClient()


@pytest.fixture(scope='module')
@pytest.fixture(scope="module")
def topic(publisher_client):
topic_path = publisher_client.topic_path(PROJECT, TOPIC)

Expand All @@ -42,16 +44,19 @@ def topic(publisher_client):

yield topic_path

publisher_client.delete_topic(topic_path)

@pytest.fixture(scope='module')

@pytest.fixture(scope="module")
def subscriber_client():
yield pubsub_v1.SubscriberClient()


@pytest.fixture
def subscription(subscriber_client, topic):
subscription_path = subscriber_client.subscription_path(
PROJECT, SUBSCRIPTION)
PROJECT, SUBSCRIPTION
)

try:
subscriber_client.delete_subscription(subscription_path)
Expand All @@ -62,6 +67,8 @@ def subscription(subscriber_client, topic):

yield subscription_path

subscriber_client.delete_subscription(subscription_path)


def test_get_topic_policy(topic, capsys):
iam.get_topic_policy(PROJECT, TOPIC)
Expand All @@ -81,16 +88,16 @@ def test_set_topic_policy(publisher_client, topic):
iam.set_topic_policy(PROJECT, TOPIC)

policy = publisher_client.get_iam_policy(topic)
assert 'roles/pubsub.publisher' in str(policy)
assert 'allUsers' in str(policy)
assert "roles/pubsub.publisher" in str(policy)
assert "allUsers" in str(policy)


def test_set_subscription_policy(subscriber_client, subscription):
iam.set_subscription_policy(PROJECT, SUBSCRIPTION)

policy = subscriber_client.get_iam_policy(subscription)
assert 'roles/pubsub.viewer' in str(policy)
assert 'allUsers' in str(policy)
assert "roles/pubsub.viewer" in str(policy)
assert "allUsers" in str(policy)


def test_check_topic_permissions(topic, capsys):
Expand All @@ -99,7 +106,7 @@ def test_check_topic_permissions(topic, capsys):
out, _ = capsys.readouterr()

assert topic in out
assert 'pubsub.topics.publish' in out
assert "pubsub.topics.publish" in out


def test_check_subscription_permissions(subscription, capsys):
Expand All @@ -108,4 +115,4 @@ def test_check_subscription_permissions(subscription, capsys):
out, _ = capsys.readouterr()

assert subscription in out
assert 'pubsub.subscriptions.consume' in out
assert "pubsub.subscriptions.consume" in out
Loading

0 comments on commit 6e46e5d

Please sign in to comment.