From d24099c40104e4786534b107b98cc9ed5672ba72 Mon Sep 17 00:00:00 2001 From: Steve Pulec Date: Sat, 29 Nov 2014 22:37:48 -0500 Subject: [PATCH] Add proper SNS paging. Closes #144, #145. --- moto/sns/models.py | 30 ++++++++++++++++------ moto/sns/responses.py | 15 ++++++----- tests/test_sns/test_subscriptions.py | 37 ++++++++++++++++++++++++++++ tests/test_sns/test_topics.py | 23 ++++++++++++++++- 4 files changed, 90 insertions(+), 15 deletions(-) diff --git a/moto/sns/models.py b/moto/sns/models.py index 60864bc21..39e1d46b2 100644 --- a/moto/sns/models.py +++ b/moto/sns/models.py @@ -1,5 +1,6 @@ from __future__ import unicode_literals +from collections import OrderedDict import datetime import uuid @@ -13,6 +14,7 @@ from moto.sqs import sqs_backends from .utils import make_arn_for_topic, make_arn_for_subscription DEFAULT_ACCOUNT_ID = 123456789012 +DEFAULT_PAGE_SIZE = 100 class Topic(object): @@ -32,7 +34,7 @@ class Topic(object): def publish(self, message): message_id = six.text_type(uuid.uuid4()) - subscriptions = self.sns_backend.list_subscriptions(self.arn) + subscriptions, _ = self.sns_backend.list_subscriptions(self.arn) for subscription in subscriptions: subscription.publish(message, message_id) return message_id @@ -77,16 +79,27 @@ class Subscription(object): class SNSBackend(BaseBackend): def __init__(self): - self.topics = {} - self.subscriptions = {} + self.topics = OrderedDict() + self.subscriptions = OrderedDict() def create_topic(self, name): topic = Topic(name, self) self.topics[topic.arn] = topic return topic - def list_topics(self): - return self.topics.values() + def _get_values_nexttoken(self, values_map, next_token=None): + if next_token is None: + next_token = 0 + next_token = int(next_token) + values = list(values_map.values())[next_token: next_token + DEFAULT_PAGE_SIZE] + if len(values) == DEFAULT_PAGE_SIZE: + next_token = next_token + DEFAULT_PAGE_SIZE + else: + next_token = None + return values, next_token + + def list_topics(self, next_token=None): + return self._get_values_nexttoken(self.topics, next_token) def delete_topic(self, arn): self.topics.pop(arn) @@ -107,12 +120,13 @@ class SNSBackend(BaseBackend): def unsubscribe(self, subscription_arn): self.subscriptions.pop(subscription_arn) - def list_subscriptions(self, topic_arn=None): + def list_subscriptions(self, topic_arn=None, next_token=None): if topic_arn: topic = self.get_topic(topic_arn) - return [sub for sub in self.subscriptions.values() if sub.topic == topic] + filtered = OrderedDict([(k, sub) for k, sub in self.subscriptions.items() if sub.topic == topic]) + return self._get_values_nexttoken(filtered, next_token) else: - return self.subscriptions.values() + return self._get_values_nexttoken(self.subscriptions, next_token) def publish(self, topic_arn, message): topic = self.get_topic(topic_arn) diff --git a/moto/sns/responses.py b/moto/sns/responses.py index 19ebdd04c..90d338f88 100644 --- a/moto/sns/responses.py +++ b/moto/sns/responses.py @@ -28,13 +28,14 @@ class SNSResponse(BaseResponse): }) def list_topics(self): - topics = self.backend.list_topics() + next_token = self._get_param('NextToken') + topics, next_token = self.backend.list_topics(next_token=next_token) return json.dumps({ 'ListTopicsResponse': { 'ListTopicsResult': { 'Topics': [{'TopicArn': topic.arn} for topic in topics], - 'NextToken': None, + 'NextToken': next_token, } }, 'ResponseMetadata': { @@ -124,7 +125,8 @@ class SNSResponse(BaseResponse): }) def list_subscriptions(self): - subscriptions = self.backend.list_subscriptions() + next_token = self._get_param('NextToken') + subscriptions, next_token = self.backend.list_subscriptions(next_token=next_token) return json.dumps({ "ListSubscriptionsResponse": { @@ -136,7 +138,7 @@ class SNSResponse(BaseResponse): "Owner": subscription.topic.account_id, "Endpoint": subscription.endpoint, } for subscription in subscriptions], - 'NextToken': None, + 'NextToken': next_token, }, "ResponseMetadata": { "RequestId": "384ac68d-3775-11df-8963-01868b7c937a", @@ -146,7 +148,8 @@ class SNSResponse(BaseResponse): def list_subscriptions_by_topic(self): topic_arn = self._get_param('TopicArn') - subscriptions = self.backend.list_subscriptions(topic_arn) + next_token = self._get_param('NextToken') + subscriptions, next_token = self.backend.list_subscriptions(topic_arn, next_token=next_token) return json.dumps({ "ListSubscriptionsByTopicResponse": { @@ -158,7 +161,7 @@ class SNSResponse(BaseResponse): "Owner": subscription.topic.account_id, "Endpoint": subscription.endpoint, } for subscription in subscriptions], - 'NextToken': None, + 'NextToken': next_token, }, "ResponseMetadata": { "RequestId": "384ac68d-3775-11df-8963-01868b7c937a", diff --git a/tests/test_sns/test_subscriptions.py b/tests/test_sns/test_subscriptions.py index 4fc14ffe5..a202edf36 100644 --- a/tests/test_sns/test_subscriptions.py +++ b/tests/test_sns/test_subscriptions.py @@ -4,6 +4,7 @@ import boto import sure # noqa from moto import mock_sns +from moto.sns.models import DEFAULT_PAGE_SIZE @mock_sns @@ -48,3 +49,39 @@ def test_getting_subscriptions_by_topic(): topic1_subscriptions = conn.get_all_subscriptions_by_topic(topic1_arn)["ListSubscriptionsByTopicResponse"]["ListSubscriptionsByTopicResult"]["Subscriptions"] topic1_subscriptions.should.have.length_of(1) topic1_subscriptions[0]['Endpoint'].should.equal("http://example1.com/") + + +@mock_sns +def test_subscription_paging(): + conn = boto.connect_sns() + conn.create_topic("topic1") + conn.create_topic("topic2") + + topics_json = conn.get_all_topics() + topics = topics_json["ListTopicsResponse"]["ListTopicsResult"]["Topics"] + topic1_arn = topics[0]['TopicArn'] + topic2_arn = topics[1]['TopicArn'] + + for index in range(DEFAULT_PAGE_SIZE + int(DEFAULT_PAGE_SIZE / 3)): + conn.subscribe(topic1_arn, 'email', 'email_' + str(index) + '@test.com') + conn.subscribe(topic2_arn, 'email', 'email_' + str(index) + '@test.com') + + all_subscriptions = conn.get_all_subscriptions() + all_subscriptions["ListSubscriptionsResponse"]["ListSubscriptionsResult"]["Subscriptions"].should.have.length_of(DEFAULT_PAGE_SIZE) + next_token = all_subscriptions["ListSubscriptionsResponse"]["ListSubscriptionsResult"]["NextToken"] + next_token.should.equal(DEFAULT_PAGE_SIZE) + + all_subscriptions = conn.get_all_subscriptions(next_token=next_token * 2) + all_subscriptions["ListSubscriptionsResponse"]["ListSubscriptionsResult"]["Subscriptions"].should.have.length_of(int(DEFAULT_PAGE_SIZE * 2 / 3)) + next_token = all_subscriptions["ListSubscriptionsResponse"]["ListSubscriptionsResult"]["NextToken"] + next_token.should.equal(None) + + topic1_subscriptions = conn.get_all_subscriptions_by_topic(topic1_arn) + topic1_subscriptions["ListSubscriptionsByTopicResponse"]["ListSubscriptionsByTopicResult"]["Subscriptions"].should.have.length_of(DEFAULT_PAGE_SIZE) + next_token = topic1_subscriptions["ListSubscriptionsByTopicResponse"]["ListSubscriptionsByTopicResult"]["NextToken"] + next_token.should.equal(DEFAULT_PAGE_SIZE) + + topic1_subscriptions = conn.get_all_subscriptions_by_topic(topic1_arn, next_token=next_token) + topic1_subscriptions["ListSubscriptionsByTopicResponse"]["ListSubscriptionsByTopicResult"]["Subscriptions"].should.have.length_of(int(DEFAULT_PAGE_SIZE / 3)) + next_token = topic1_subscriptions["ListSubscriptionsByTopicResponse"]["ListSubscriptionsByTopicResult"]["NextToken"] + next_token.should.equal(None) diff --git a/tests/test_sns/test_topics.py b/tests/test_sns/test_topics.py index 427c8c003..817426244 100644 --- a/tests/test_sns/test_topics.py +++ b/tests/test_sns/test_topics.py @@ -5,7 +5,7 @@ import six import sure # noqa from moto import mock_sns -from moto.sns.models import DEFAULT_TOPIC_POLICY, DEFAULT_EFFECTIVE_DELIVERY_POLICY +from moto.sns.models import DEFAULT_TOPIC_POLICY, DEFAULT_EFFECTIVE_DELIVERY_POLICY, DEFAULT_PAGE_SIZE @mock_sns @@ -77,3 +77,24 @@ def test_topic_attributes(): attributes["Policy"].should.equal("{'foo': 'bar'}") attributes["DisplayName"].should.equal("My display name") attributes["DeliveryPolicy"].should.equal("{'http': {'defaultHealthyRetryPolicy': {'numRetries': 5}}}") + + +@mock_sns +def test_topic_paging(): + conn = boto.connect_sns() + for index in range(DEFAULT_PAGE_SIZE + int(DEFAULT_PAGE_SIZE / 2)): + conn.create_topic("some-topic_" + str(index)) + + topics_json = conn.get_all_topics() + topics_list = topics_json["ListTopicsResponse"]["ListTopicsResult"]["Topics"] + next_token = topics_json["ListTopicsResponse"]["ListTopicsResult"]["NextToken"] + + len(topics_list).should.equal(DEFAULT_PAGE_SIZE) + next_token.should.equal(DEFAULT_PAGE_SIZE) + + topics_json = conn.get_all_topics(next_token=next_token) + topics_list = topics_json["ListTopicsResponse"]["ListTopicsResult"]["Topics"] + next_token = topics_json["ListTopicsResponse"]["ListTopicsResult"]["NextToken"] + + topics_list.should.have.length_of(int(DEFAULT_PAGE_SIZE / 2)) + next_token.should.equal(None)