Events: Fix pagination for list_rules/list_rule_names_by_target (#3781)

moto/events/models.py

@@ -12,6 +12,7 @@ from operator import lt, le, eq, ge, gt
 
 from boto3 import Session
 
+from collections import OrderedDict
 from moto.core.exceptions import JsonRESTError
 from moto.core import ACCOUNT_ID, BaseBackend, CloudFormationModel, BaseModel
 from moto.core.utils import unix_time, iso_8601_datetime_without_milliseconds
@@ -22,10 +23,13 @@ from moto.events.exceptions import (
     InvalidEventPatternException,
     IllegalStatusException,
 )
+from moto.utilities.paginator import paginate
 from moto.utilities.tagging_service import TaggingService
 
 from uuid import uuid4
 
+from .utils import PAGINATION_MODEL
+
 
 class Rule(CloudFormationModel):
     Arn = namedtuple("Arn", ["service", "resource_type", "resource_id"])
@@ -901,10 +905,7 @@ class EventsBackend(BaseBackend):
     _RATE_REGEX = re.compile(r"^rate\(\d*\s(minute|minutes|hour|hours|day|days)\)")
 
     def __init__(self, region_name):
-        self.rules = {}
+        self.rules = OrderedDict()
-        # This array tracks the order in which the rules have been added, since
-        # 2.6 doesn't have OrderedDicts.
-        self.rules_order = []
         self.next_tokens = {}
         self.region_name = region_name
         self.event_buses = {}
@@ -932,9 +933,6 @@ class EventsBackend(BaseBackend):
     def _add_default_event_bus(self):
         self.event_buses["default"] = EventBus(self.region_name, "default")
 
-    def _get_rule_by_index(self, i):
-        return self.rules.get(self.rules_order[i])
-
     def _gen_next_token(self, index):
         token = os.urandom(128).encode("base64")
         self.next_tokens[token] = index
@@ -1022,7 +1020,6 @@ class EventsBackend(BaseBackend):
             targets=targets,
         )
         self.rules[name] = rule
-        self.rules_order.append(name)
 
         if tags:
             self.tagger.tag_resource(rule.arn, tags)
@@ -1030,7 +1027,6 @@ class EventsBackend(BaseBackend):
         return rule
 
     def delete_rule(self, name):
-        self.rules_order.pop(self.rules_order.index(name))
         arn = self.rules.get(name).arn
         if self.tagger.has_tags(arn):
             self.tagger.delete_all_tags_for_resource(arn)
@@ -1056,26 +1052,18 @@ class EventsBackend(BaseBackend):
 
         return False
 
+    @paginate(pagination_model=PAGINATION_MODEL)
     def list_rule_names_by_target(self, target_arn, next_token=None, limit=None):
         matching_rules = []
-        return_obj = {}
 
-        start_index, end_index, new_next_token = self._process_token_and_limits(
+        for _, rule in self.rules.items():
-            len(self.rules), next_token, limit
-        )
-
-        for i in range(start_index, end_index):
-            rule = self._get_rule_by_index(i)
             for target in rule.targets:
                 if target["Arn"] == target_arn:
-                    matching_rules.append(rule.name)
+                    matching_rules.append(rule)
 
-        return_obj["RuleNames"] = matching_rules
+        return matching_rules
-        if new_next_token is not None:
-            return_obj["NextToken"] = new_next_token
-
-        return return_obj
 
+    @paginate(pagination_model=PAGINATION_MODEL)
     def list_rules(self, prefix=None, next_token=None, limit=None):
         match_string = ".*"
         if prefix is not None:
@@ -1084,22 +1072,12 @@ class EventsBackend(BaseBackend):
         match_regex = re.compile(match_string)
 
         matching_rules = []
-        return_obj = {}
 
-        start_index, end_index, new_next_token = self._process_token_and_limits(
+        for name, rule in self.rules.items():
-            len(self.rules), next_token, limit
+            if match_regex.match(name):
-        )
-
-        for i in range(start_index, end_index):
-            rule = self._get_rule_by_index(i)
-            if match_regex.match(rule.name):
                 matching_rules.append(rule)
 
-        return_obj["Rules"] = matching_rules
+        return matching_rules
-        if new_next_token is not None:
-            return_obj["NextToken"] = new_next_token
-
-        return return_obj
 
     def list_targets_by_rule(self, rule, next_token=None, limit=None):
         # We'll let a KeyError exception be thrown for response to handle if

moto/events/responses.py

@@ -126,25 +126,26 @@ class EventsHandler(BaseResponse):
         if not target_arn:
             return self.error("ValidationException", "Parameter TargetArn is required.")
 
-        rule_names = self.events_backend.list_rule_names_by_target(
+        rules, token = self.events_backend.list_rule_names_by_target(
-            target_arn, next_token, limit
+            target_arn=target_arn, next_token=next_token, limit=limit
         )
 
-        return json.dumps(rule_names), self.response_headers
+        res = {"RuleNames": [rule.name for rule in rules], "NextToken": token}
+
+        return json.dumps(res), self.response_headers
 
     def list_rules(self):
         prefix = self._get_param("NamePrefix")
         next_token = self._get_param("NextToken")
         limit = self._get_param("Limit")
 
-        rules = self.events_backend.list_rules(prefix, next_token, limit)
+        rules, token = self.events_backend.list_rules(
-        rules_obj = {"Rules": []}
+            prefix=prefix, next_token=next_token, limit=limit
-
+        )
-        for rule in rules["Rules"]:
+        rules_obj = {
-            rules_obj["Rules"].append(rule.describe())
+            "Rules": [rule.describe() for rule in rules],
-
+            "NextToken": token,
-        if rules.get("NextToken"):
+        }
-            rules_obj["NextToken"] = rules["NextToken"]
 
         return json.dumps(rules_obj), self.response_headers
 

moto/events/utils.py

@@ -0,0 +1,16 @@
+PAGINATION_MODEL = {
+    "list_rules": {
+        "input_token": "next_token",
+        "limit_key": "limit",
+        "limit_default": 50,
+        "page_ending_range_keys": ["arn"],
+        "fail_on_invalid_token": False,
+    },
+    "list_rule_names_by_target": {
+        "input_token": "next_token",
+        "limit_key": "limit",
+        "limit_default": 50,
+        "page_ending_range_keys": ["arn"],
+        "fail_on_invalid_token": False,
+    },
+}

moto/kinesis/utils.py

@@ -18,6 +18,17 @@ PAGINATION_MODEL = {
 }
 
 
+PAGINATION_MODEL = {
+    "list_shards": {
+        "input_token": "next_token",
+        "limit_key": "limit",
+        "limit_default": 10000,
+        "page_ending_range_keys": ["ShardId"],
+        "fail_on_invalid_token": False,
+    },
+}
+
+
 def compose_new_shard_iterator(
     stream_name, shard, shard_iterator_type, starting_sequence_number, at_timestamp
 ):

tests/test_events/test_events.py

@@ -134,9 +134,46 @@ def test_put_rule_error_schedule_expression_custom_event_bus():
 def test_list_rules():
     client = generate_environment()
     response = client.list_rules()
+    rules = response["Rules"]
+    rules.should.have.length_of(len(RULES))
 
-    assert response is not None
+
-    assert len(response["Rules"]) > 0
+@mock_events
+def test_list_rules_with_token():
+    client = generate_environment()
+    response = client.list_rules()
+    response.shouldnt.have.key("NextToken")
+    rules = response["Rules"]
+    rules.should.have.length_of(len(RULES))
+    #
+    response = client.list_rules(Limit=1)
+    response.should.have.key("NextToken")
+    rules = response["Rules"]
+    rules.should.have.length_of(1)
+    #
+    response = client.list_rules(NextToken=response["NextToken"])
+    response.shouldnt.have.key("NextToken")
+    rules = response["Rules"]
+    rules.should.have.length_of(2)
+
+
+@mock_events
+def test_list_rules_with_prefix_and_token():
+    client = generate_environment()
+    response = client.list_rules(NamePrefix="test")
+    response.shouldnt.have.key("NextToken")
+    rules = response["Rules"]
+    rules.should.have.length_of(len(RULES))
+    #
+    response = client.list_rules(NamePrefix="test", Limit=1)
+    response.should.have.key("NextToken")
+    rules = response["Rules"]
+    rules.should.have.length_of(1)
+    #
+    response = client.list_rules(NamePrefix="test", NextToken=response["NextToken"])
+    response.shouldnt.have.key("NextToken")
+    rules = response["Rules"]
+    rules.should.have.length_of(2)
 
 
 @mock_events
@@ -242,6 +279,23 @@ def test_list_rule_names_by_target():
         assert rule in test_2_target["Rules"]
 
 
+@mock_events
+def test_list_rule_names_by_target_using_limit():
+    test_1_target = TARGETS["test-target-1"]
+    client = generate_environment()
+
+    response = client.list_rule_names_by_target(TargetArn=test_1_target["Arn"], Limit=1)
+    print(response)
+    response.should.have.key("NextToken")
+    response["RuleNames"].should.have.length_of(1)
+    #
+    response = client.list_rule_names_by_target(
+        TargetArn=test_1_target["Arn"], NextToken=response["NextToken"]
+    )
+    response.shouldnt.have.key("NextToken")
+    response["RuleNames"].should.have.length_of(1)
+
+
 @mock_events
 def test_delete_rule():
     client = generate_environment()