Skip to content

Commit

Permalink
test: put RBAC tests into their own suite (#51)
Browse files Browse the repository at this point in the history
Co-authored-by: Prad Nelluru <pradn@google.com>
  • Loading branch information
plamut and pradn committed Mar 24, 2020
1 parent 9abb047 commit a3ee941
Showing 1 changed file with 32 additions and 47 deletions.
79 changes: 32 additions & 47 deletions tests/system.py
Expand Up @@ -500,7 +500,7 @@ def test_streaming_pull_ack_deadline(
)

# publish some messages and wait for completion
self._publish_messages(publisher, topic_path, batch_sizes=[2])
_publish_messages(publisher, topic_path, batch_sizes=[2])

# subscribe to the topic
callback = StreamingPullCallback(
Expand Down Expand Up @@ -543,7 +543,7 @@ def test_streaming_pull_max_messages(
subscriber.create_subscription(subscription_path, topic_path)

batch_sizes = (7, 4, 8, 2, 10, 1, 3, 8, 6, 1) # total: 50
self._publish_messages(publisher, topic_path, batch_sizes=batch_sizes)
_publish_messages(publisher, topic_path, batch_sizes=batch_sizes)

# now subscribe and do the main part, check for max pending messages
total_messages = sum(batch_sizes)
Expand Down Expand Up @@ -585,10 +585,12 @@ def test_streaming_pull_max_messages(
finally:
subscription_future.cancel() # trigger clean shutdown

@pytest.mark.skipif(
"KOKORO_GFILE_DIR" not in os.environ,
reason="Requires Kokoro environment with a service account with limited role.",
)

@pytest.mark.skipif(
"KOKORO_GFILE_DIR" not in os.environ,
reason="Requires Kokoro environment with a service account with limited role.",
)
class TestBasicRBAC(object):
def test_streaming_pull_subscriber_permissions_sufficient(
self, publisher, topic_path, subscriber, subscription_path, cleanup
):
Expand All @@ -611,7 +613,7 @@ def test_streaming_pull_subscriber_permissions_sufficient(
# successfully pulls and processes it.
callback = StreamingPullCallback(processing_time=0.01, resolve_at_msg_count=1)
future = streaming_pull_subscriber.subscribe(subscription_path, callback)
self._publish_messages(publisher, topic_path, batch_sizes=[1])
_publish_messages(publisher, topic_path, batch_sizes=[1])

try:
callback.done_future.result(timeout=10)
Expand All @@ -624,10 +626,6 @@ def test_streaming_pull_subscriber_permissions_sufficient(
finally:
future.cancel()

@pytest.mark.skipif(
"KOKORO_GFILE_DIR" not in os.environ,
reason="Requires Kokoro environment with a service account with limited role.",
)
def test_publisher_role_can_publish_messages(
self, publisher, topic_path, subscriber, subscription_path, cleanup
):
Expand All @@ -646,18 +644,14 @@ def test_publisher_role_can_publish_messages(
)
publisher_only_client = type(publisher).from_service_account_file(filename)

self._publish_messages(publisher_only_client, topic_path, batch_sizes=[2])
_publish_messages(publisher_only_client, topic_path, batch_sizes=[2])

response = subscriber.pull(subscription_path, max_messages=2)
assert len(response.received_messages) == 2

@pytest.mark.skip(
"Snapshot creation is not instant on the backend, causing test falkiness."
)
@pytest.mark.skipif(
"KOKORO_GFILE_DIR" not in os.environ,
reason="Requires Kokoro environment with a service account with limited role.",
)
def test_snapshot_seek_subscriber_permissions_sufficient(
self, project, publisher, topic_path, subscriber, subscription_path, cleanup
):
Expand All @@ -682,13 +676,13 @@ def test_snapshot_seek_subscriber_permissions_sufficient(
subscriber_only_client = type(subscriber).from_service_account_file(filename)

# Publish two messages and create a snapshot inbetween.
self._publish_messages(publisher, topic_path, batch_sizes=[1])
_publish_messages(publisher, topic_path, batch_sizes=[1])
response = subscriber.pull(subscription_path, max_messages=10)
assert len(response.received_messages) == 1

subscriber.create_snapshot(snapshot_path, subscription_path)

self._publish_messages(publisher, topic_path, batch_sizes=[1])
_publish_messages(publisher, topic_path, batch_sizes=[1])
response = subscriber.pull(subscription_path, max_messages=10)
assert len(response.received_messages) == 1

Expand All @@ -699,10 +693,6 @@ def test_snapshot_seek_subscriber_permissions_sufficient(
response = subscriber.pull(subscription_path, max_messages=10)
assert len(response.received_messages) == 1

@pytest.mark.skipif(
"KOKORO_GFILE_DIR" not in os.environ,
reason="Requires Kokoro environment with a service account with limited role.",
)
def test_viewer_role_can_list_resources(
self, project, publisher, topic_path, subscriber, cleanup
):
Expand All @@ -727,10 +717,6 @@ def test_viewer_role_can_list_resources(
next(iter(viewer_only_subscriber.list_subscriptions(project_path)), None)
next(iter(viewer_only_subscriber.list_snapshots(project_path)), None)

@pytest.mark.skipif(
"KOKORO_GFILE_DIR" not in os.environ,
reason="Requires Kokoro environment with a service account with limited role.",
)
def test_editor_role_can_create_resources(
self, project, publisher, topic_path, subscriber, subscription_path, cleanup
):
Expand All @@ -754,30 +740,29 @@ def test_editor_role_can_create_resources(
editor_subscriber.create_subscription(subscription_path, topic_path)
editor_subscriber.create_snapshot(snapshot_path, subscription_path)

def _publish_messages(self, publisher, topic_path, batch_sizes):
"""Publish ``count`` messages in batches and wait until completion."""
publish_futures = []
msg_counter = itertools.count(start=1)

for batch_size in batch_sizes:
msg_batch = self._make_messages(count=batch_size)
for msg in msg_batch:
future = publisher.publish(
topic_path, msg, seq_num=str(next(msg_counter))
)
publish_futures.append(future)
time.sleep(0.1)
def _publish_messages(publisher, topic_path, batch_sizes):
"""Publish ``count`` messages in batches and wait until completion."""
publish_futures = []
msg_counter = itertools.count(start=1)

# wait untill all messages have been successfully published
for future in publish_futures:
future.result(timeout=30)
for batch_size in batch_sizes:
msg_batch = _make_messages(count=batch_size)
for msg in msg_batch:
future = publisher.publish(topic_path, msg, seq_num=str(next(msg_counter)))
publish_futures.append(future)
time.sleep(0.1)

def _make_messages(self, count):
messages = [
u"message {}/{}".format(i, count).encode("utf-8")
for i in range(1, count + 1)
]
return messages
# wait untill all messages have been successfully published
for future in publish_futures:
future.result(timeout=30)


def _make_messages(count):
messages = [
u"message {}/{}".format(i, count).encode("utf-8") for i in range(1, count + 1)
]
return messages


class AckCallback(object):
Expand Down

0 comments on commit a3ee941

Please sign in to comment.