From 4befb671f0a686e27d57809af042895b269d5e1d Mon Sep 17 00:00:00 2001 From: Bert Blommers Date: Fri, 3 Dec 2021 15:33:35 -0100 Subject: [PATCH] Paginator - make development easier + docs (#4650) --- .../development_tips/utilities.rst | 29 +- moto/cognitoidp/models.py | 8 +- moto/cognitoidp/utils.py | 8 +- moto/ds/models.py | 8 +- moto/ds/utils.py | 4 +- moto/events/models.py | 4 +- moto/events/utils.py | 4 +- moto/glue/models.py | 2 +- moto/iot/utils.py | 2 +- moto/kinesis/models.py | 2 +- moto/kinesis/utils.py | 4 +- moto/logs/models.py | 4 +- moto/logs/utils.py | 4 +- moto/route53/models.py | 4 +- moto/route53/utils.py | 2 +- moto/route53resolver/models.py | 20 +- moto/route53resolver/responses.py | 23 +- moto/route53resolver/utils.py | 14 +- moto/stepfunctions/utils.py | 4 +- moto/utilities/paginator.py | 92 ++++-- tests/test_utilities/test_paginator.py | 299 ++++++++++++++++++ 21 files changed, 439 insertions(+), 102 deletions(-) create mode 100644 tests/test_utilities/test_paginator.py diff --git a/docs/docs/contributing/development_tips/utilities.rst b/docs/docs/contributing/development_tips/utilities.rst index 579e89417..ad6157172 100644 --- a/docs/docs/contributing/development_tips/utilities.rst +++ b/docs/docs/contributing/development_tips/utilities.rst @@ -59,10 +59,21 @@ See the below example how it works in practice: # The default limit of the above parameter is not provided "limit_default": 100, # - # The collection of keys/attributes that should guarantee uniqueness for a given resource. + # One or more attributes that guarantee uniqueness for a given resource. # For most resources it will just be an ID, or ARN, which is always unique. - # In order to know what is the last resource that we're sending back, an encoded version of these attributes is used as the NextToken. - "page_ending_range_keys": ["start_date", "execution_arn"], + # An encoded version of these attributes is used as the NextToken. + "unique_attribute": "arn", + # Provide a list if only a combination of attributes is guaranteed to be unique + "unique_attribute": ["start_date", "execution_arn"], + # + # By default, an exception will be thrown if the user-provided next_token is invalid + "fail_on_invalid_token": True # Default value - no need to specify this + # This can be customized to: + # - silently fail, and just return an empty list + "fail_on_invalid_token": False, + # - throw a custom exception, by providing an exception class + # The paginator will `raise CustomException()` or `raise CustomException(invalid_token)` + "fail_on_invalid_token": CustomException }, # another method that will use different configuration options "list_other_things": { @@ -72,9 +83,17 @@ See the below example how it works in practice: # The decorator with the pagination logic @paginate(pagination_model=PAGINATION_MODEL) - # Note that this method does not list the 'next_token'/'max_results'-arguments - # The decorator uses them, but our logic doesn't need them + # Note that this method does not have the 'next_token'/'max_results'-arguments def list_resources(self): # Note that we simply return all resources - the decorator takes care of all pagination magic return self.full_list_of_resources + @paginate(pagination_model=PAGINATION_MODEL) + # If we do need the 'next_token'/'max_results'-arguments, just add them to the function + # The decorator will only pass them along if required + def list_other_things(self, max_results=None): + if max_results == "42": + # Custom validation magic + pass + return self.full_list_of_resources + diff --git a/moto/cognitoidp/models.py b/moto/cognitoidp/models.py index 6463460f3..38dbbfc2f 100644 --- a/moto/cognitoidp/models.py +++ b/moto/cognitoidp/models.py @@ -767,7 +767,7 @@ class CognitoIdpBackend(BaseBackend): } @paginate(pagination_model=PAGINATION_MODEL) - def list_user_pools(self, max_results=None, next_token=None): + def list_user_pools(self): return list(self.user_pools.values()) def describe_user_pool(self, user_pool_id): @@ -827,7 +827,7 @@ class CognitoIdpBackend(BaseBackend): return user_pool_client @paginate(pagination_model=PAGINATION_MODEL) - def list_user_pool_clients(self, user_pool_id, max_results=None, next_token=None): + def list_user_pool_clients(self, user_pool_id): user_pool = self.describe_user_pool(user_pool_id) return list(user_pool.clients.values()) @@ -868,7 +868,7 @@ class CognitoIdpBackend(BaseBackend): return identity_provider @paginate(pagination_model=PAGINATION_MODEL) - def list_identity_providers(self, user_pool_id, max_results=None, next_token=None): + def list_identity_providers(self, user_pool_id): user_pool = self.describe_user_pool(user_pool_id) return list(user_pool.identity_providers.values()) @@ -1071,7 +1071,7 @@ class CognitoIdpBackend(BaseBackend): raise NotAuthorizedError("Invalid token") @paginate(pagination_model=PAGINATION_MODEL) - def list_users(self, user_pool_id, pagination_token=None, limit=None): + def list_users(self, user_pool_id): user_pool = self.describe_user_pool(user_pool_id) return list(user_pool.users.values()) diff --git a/moto/cognitoidp/utils.py b/moto/cognitoidp/utils.py index be00dbbbd..8a4970957 100644 --- a/moto/cognitoidp/utils.py +++ b/moto/cognitoidp/utils.py @@ -16,25 +16,25 @@ PAGINATION_MODEL = { "input_token": "next_token", "limit_key": "max_results", "limit_default": 60, - "page_ending_range_keys": ["arn"], + "unique_attribute": "arn", }, "list_user_pool_clients": { "input_token": "next_token", "limit_key": "max_results", "limit_default": 60, - "page_ending_range_keys": ["id"], + "unique_attribute": "id", }, "list_identity_providers": { "input_token": "next_token", "limit_key": "max_results", "limit_default": 60, - "page_ending_range_keys": ["name"], + "unique_attribute": "name", }, "list_users": { "input_token": "pagination_token", "limit_key": "limit", "limit_default": 60, - "page_ending_range_keys": ["id"], + "unique_attribute": "id", }, } diff --git a/moto/ds/models.py b/moto/ds/models.py index c03a2f8ea..446f150e4 100644 --- a/moto/ds/models.py +++ b/moto/ds/models.py @@ -452,9 +452,7 @@ class DirectoryServiceBackend(BaseBackend): directory.enable_sso(True) @paginate(pagination_model=PAGINATION_MODEL) - def describe_directories( - self, directory_ids=None, next_token=None, limit=0 - ): # pylint: disable=unused-argument + def describe_directories(self, directory_ids=None): """Return info on all directories or directories with matching IDs.""" for directory_id in directory_ids or self.directories: self._validate_directory_id(directory_id) @@ -506,9 +504,7 @@ class DirectoryServiceBackend(BaseBackend): self.tagger.untag_resource_using_names(resource_id, tag_keys) @paginate(pagination_model=PAGINATION_MODEL) - def list_tags_for_resource( - self, resource_id, next_token=None, limit=None, - ): # pylint: disable=unused-argument + def list_tags_for_resource(self, resource_id): """List all tags on a directory.""" self._validate_directory_id(resource_id) return self.tagger.list_tags_for_resource(resource_id).get("Tags") diff --git a/moto/ds/utils.py b/moto/ds/utils.py index f994da4c4..f297d393b 100644 --- a/moto/ds/utils.py +++ b/moto/ds/utils.py @@ -5,12 +5,12 @@ PAGINATION_MODEL = { "input_token": "next_token", "limit_key": "limit", "limit_default": 100, # This should be the sum of the directory limits - "page_ending_range_keys": ["directory_id"], + "unique_attribute": "directory_id", }, "list_tags_for_resource": { "input_token": "next_token", "limit_key": "limit", "limit_default": 50, - "page_ending_range_keys": ["Key"], + "unique_attribute": "Key", }, } diff --git a/moto/events/models.py b/moto/events/models.py index 676476313..85db1f455 100644 --- a/moto/events/models.py +++ b/moto/events/models.py @@ -1070,7 +1070,7 @@ class EventsBackend(BaseBackend): return False @paginate(pagination_model=PAGINATION_MODEL) - def list_rule_names_by_target(self, target_arn, next_token=None, limit=None): + def list_rule_names_by_target(self, target_arn): matching_rules = [] for _, rule in self.rules.items(): @@ -1081,7 +1081,7 @@ class EventsBackend(BaseBackend): return matching_rules @paginate(pagination_model=PAGINATION_MODEL) - def list_rules(self, prefix=None, next_token=None, limit=None): + def list_rules(self, prefix=None): match_string = ".*" if prefix is not None: match_string = "^" + prefix + match_string diff --git a/moto/events/utils.py b/moto/events/utils.py index a333054b6..ad359ab6d 100644 --- a/moto/events/utils.py +++ b/moto/events/utils.py @@ -3,14 +3,14 @@ PAGINATION_MODEL = { "input_token": "next_token", "limit_key": "limit", "limit_default": 50, - "page_ending_range_keys": ["arn"], + "unique_attribute": "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"], + "unique_attribute": "arn", "fail_on_invalid_token": False, }, } diff --git a/moto/glue/models.py b/moto/glue/models.py index 2ca2a3315..616d6312c 100644 --- a/moto/glue/models.py +++ b/moto/glue/models.py @@ -25,7 +25,7 @@ class GlueBackend(BaseBackend): "input_token": "next_token", "limit_key": "max_results", "limit_default": 100, - "page_ending_range_keys": ["name"], + "unique_attribute": "name", }, } diff --git a/moto/iot/utils.py b/moto/iot/utils.py index c4ea599ff..1a08f56b0 100644 --- a/moto/iot/utils.py +++ b/moto/iot/utils.py @@ -3,6 +3,6 @@ PAGINATION_MODEL = { "input_token": "next_token", "limit_key": "max_results", "limit_default": 100, - "page_ending_range_keys": ["jobId"], + "unique_attribute": "jobId", } } diff --git a/moto/kinesis/models.py b/moto/kinesis/models.py index 99c3f1052..5f16fb69f 100644 --- a/moto/kinesis/models.py +++ b/moto/kinesis/models.py @@ -543,7 +543,7 @@ class KinesisBackend(BaseBackend): ) @paginate(pagination_model=PAGINATION_MODEL) - def list_shards(self, stream_name, limit=None, next_token=None): + def list_shards(self, stream_name): stream = self.describe_stream(stream_name) shards = sorted(stream.shards.values(), key=lambda x: x.shard_id) return [shard.to_json() for shard in shards] diff --git a/moto/kinesis/utils.py b/moto/kinesis/utils.py index bf750cef9..9de8052a1 100644 --- a/moto/kinesis/utils.py +++ b/moto/kinesis/utils.py @@ -12,7 +12,7 @@ PAGINATION_MODEL = { "input_token": "next_token", "limit_key": "limit", "limit_default": 10000, - "page_ending_range_keys": ["ShardId"], + "unique_attribute": "ShardId", "fail_on_invalid_token": False, }, } @@ -23,7 +23,7 @@ PAGINATION_MODEL = { "input_token": "next_token", "limit_key": "limit", "limit_default": 10000, - "page_ending_range_keys": ["ShardId"], + "unique_attribute": "ShardId", "fail_on_invalid_token": False, }, } diff --git a/moto/logs/models.py b/moto/logs/models.py index 337e981d0..432009978 100644 --- a/moto/logs/models.py +++ b/moto/logs/models.py @@ -672,9 +672,7 @@ class LogsBackend(BaseBackend): del self.groups[log_group_name] @paginate(pagination_model=PAGINATION_MODEL) - def describe_log_groups( - self, log_group_name_prefix=None, limit=None, next_token=None - ): + def describe_log_groups(self, log_group_name_prefix=None): if log_group_name_prefix is None: log_group_name_prefix = "" diff --git a/moto/logs/utils.py b/moto/logs/utils.py index adc25cb02..74d832883 100644 --- a/moto/logs/utils.py +++ b/moto/logs/utils.py @@ -3,13 +3,13 @@ PAGINATION_MODEL = { "input_token": "next_token", "limit_key": "limit", "limit_default": 50, - "page_ending_range_keys": ["arn"], + "unique_attribute": "arn", "fail_on_invalid_token": False, }, "describe_log_streams": { "input_token": "next_token", "limit_key": "limit", "limit_default": 50, - "page_ending_range_keys": ["arn"], + "unique_attribute": "arn", }, } diff --git a/moto/route53/models.py b/moto/route53/models.py index 889763daf..0e694b74e 100644 --- a/moto/route53/models.py +++ b/moto/route53/models.py @@ -598,9 +598,7 @@ class Route53Backend(BaseBackend): return self.query_logging_configs[query_logging_config_id] @paginate(pagination_model=PAGINATION_MODEL) - def list_query_logging_configs( - self, hosted_zone_id=None, next_token=None, max_results=None, - ): # pylint: disable=unused-argument + def list_query_logging_configs(self, hosted_zone_id=None): """Return a list of query logging configs.""" if hosted_zone_id: # Does the hosted_zone_id exist? diff --git a/moto/route53/utils.py b/moto/route53/utils.py index 858f48878..8c0f38d27 100644 --- a/moto/route53/utils.py +++ b/moto/route53/utils.py @@ -5,6 +5,6 @@ PAGINATION_MODEL = { "input_token": "next_token", "limit_key": "max_results", "limit_default": 100, - "page_ending_range_keys": ["hosted_zone_id"], + "unique_attribute": "hosted_zone_id", }, } diff --git a/moto/route53resolver/models.py b/moto/route53resolver/models.py index 190837c74..1196fa95a 100644 --- a/moto/route53resolver/models.py +++ b/moto/route53resolver/models.py @@ -690,9 +690,7 @@ class Route53ResolverBackend(BaseBackend): return self.resolver_rule_associations[resolver_rule_association_id] @paginate(pagination_model=PAGINATION_MODEL) - def list_resolver_endpoint_ip_addresses( - self, resolver_endpoint_id, next_token=None, max_results=None, - ): # pylint: disable=unused-argument + def list_resolver_endpoint_ip_addresses(self, resolver_endpoint_id): """List IP endresses for specified resolver endpoint.""" self._validate_resolver_endpoint_id(resolver_endpoint_id) endpoint = self.resolver_endpoints[resolver_endpoint_id] @@ -752,9 +750,7 @@ class Route53ResolverBackend(BaseBackend): return True @paginate(pagination_model=PAGINATION_MODEL) - def list_resolver_endpoints( - self, filters, next_token=None, max_results=None, - ): # pylint: disable=unused-argument + def list_resolver_endpoints(self, filters): """List all resolver endpoints, using filters if specified.""" if not filters: filters = [] @@ -769,9 +765,7 @@ class Route53ResolverBackend(BaseBackend): return endpoints @paginate(pagination_model=PAGINATION_MODEL) - def list_resolver_rules( - self, filters, next_token=None, max_results=None, - ): # pylint: disable=unused-argument + def list_resolver_rules(self, filters): """List all resolver rules, using filters if specified.""" if not filters: filters = [] @@ -786,9 +780,7 @@ class Route53ResolverBackend(BaseBackend): return rules @paginate(pagination_model=PAGINATION_MODEL) - def list_resolver_rule_associations( - self, filters, next_token=None, max_results=None, - ): # pylint: disable=unused-argument + def list_resolver_rule_associations(self, filters): """List all resolver rule associations, using filters if specified.""" if not filters: filters = [] @@ -817,9 +809,7 @@ class Route53ResolverBackend(BaseBackend): ) @paginate(pagination_model=PAGINATION_MODEL) - def list_tags_for_resource( - self, resource_arn, next_token=None, max_results=None, - ): # pylint: disable=unused-argument + def list_tags_for_resource(self, resource_arn): """List all tags for the given resource.""" self._matched_arn(resource_arn) return self.tagger.list_tags_for_resource(resource_arn).get("Tags") diff --git a/moto/route53resolver/responses.py b/moto/route53resolver/responses.py index 7aee42ee6..9d51dac3f 100644 --- a/moto/route53resolver/responses.py +++ b/moto/route53resolver/responses.py @@ -156,15 +156,9 @@ class Route53ResolverResponse(BaseResponse): next_token = self._get_param("NextToken") max_results = self._get_param("MaxResults", 10) validate_args([("maxResults", max_results)]) - try: - ( - endpoints, - next_token, - ) = self.route53resolver_backend.list_resolver_endpoints( - filters, next_token=next_token, max_results=max_results - ) - except InvalidToken as exc: - raise InvalidNextTokenException() from exc + endpoints, next_token = self.route53resolver_backend.list_resolver_endpoints( + filters, next_token=next_token, max_results=max_results + ) response = { "ResolverEndpoints": [x.description() for x in endpoints], @@ -224,14 +218,9 @@ class Route53ResolverResponse(BaseResponse): resource_arn = self._get_param("ResourceArn") next_token = self._get_param("NextToken") max_results = self._get_param("MaxResults") - try: - (tags, next_token) = self.route53resolver_backend.list_tags_for_resource( - resource_arn=resource_arn, - next_token=next_token, - max_results=max_results, - ) - except InvalidToken as exc: - raise InvalidNextTokenException() from exc + tags, next_token = self.route53resolver_backend.list_tags_for_resource( + resource_arn=resource_arn, next_token=next_token, max_results=max_results, + ) response = {"Tags": tags} if next_token: diff --git a/moto/route53resolver/utils.py b/moto/route53resolver/utils.py index 0a4df012e..be13f2033 100644 --- a/moto/route53resolver/utils.py +++ b/moto/route53resolver/utils.py @@ -1,34 +1,38 @@ """Pagination control model for Route53Resolver.""" +from .exceptions import InvalidNextTokenException + PAGINATION_MODEL = { "list_resolver_endpoint_ip_addresses": { "input_token": "next_token", "limit_key": "max_results", "limit_default": 100, - "page_ending_range_keys": ["IpId"], + "unique_attribute": "IpId", }, "list_resolver_endpoints": { "input_token": "next_token", "limit_key": "max_results", "limit_default": 100, - "page_ending_range_keys": ["id"], + "unique_attribute": "id", + "fail_on_invalid_token": InvalidNextTokenException, }, "list_resolver_rules": { "input_token": "next_token", "limit_key": "max_results", "limit_default": 100, - "page_ending_range_keys": ["id"], + "unique_attribute": "id", }, "list_resolver_rule_associations": { "input_token": "next_token", "limit_key": "max_results", "limit_default": 100, - "page_ending_range_keys": ["id"], + "unique_attribute": "id", }, "list_tags_for_resource": { "input_token": "next_token", "limit_key": "max_results", "limit_default": 100, - "page_ending_range_keys": ["Key"], + "unique_attribute": "Key", + "fail_on_invalid_token": InvalidNextTokenException, }, } diff --git a/moto/stepfunctions/utils.py b/moto/stepfunctions/utils.py index b4e3f4265..20881771f 100644 --- a/moto/stepfunctions/utils.py +++ b/moto/stepfunctions/utils.py @@ -3,13 +3,13 @@ PAGINATION_MODEL = { "input_token": "next_token", "limit_key": "max_results", "limit_default": 100, - "page_ending_range_keys": ["start_date", "execution_arn"], + "unique_attribute": ["start_date", "execution_arn"], }, "list_state_machines": { "input_token": "next_token", "limit_key": "max_results", "limit_default": 100, - "page_ending_range_keys": ["creation_date", "arn"], + "unique_attribute": ["creation_date", "arn"], }, } diff --git a/moto/utilities/paginator.py b/moto/utilities/paginator.py index efdddd942..ba52ccfa6 100644 --- a/moto/utilities/paginator.py +++ b/moto/utilities/paginator.py @@ -1,4 +1,7 @@ -from functools import wraps, reduce +import inspect + +from copy import deepcopy +from functools import wraps from botocore.paginate import TokenDecoder, TokenEncoder @@ -17,19 +20,39 @@ def paginate(pagination_model, original_function=None): raise ValueError( "No pagination config for backend method: {}".format(method) ) - # We pop the pagination arguments, so the remaining kwargs (if any) - # can be used to compute the optional parameters checksum. - input_token = kwargs.pop(pagination_config.get("input_token"), None) - limit = kwargs.pop(pagination_config.get("limit_key"), None) + # Get the pagination arguments, to be used by the paginator + next_token_name = pagination_config.get("input_token", "next_token") + limit_name = pagination_config.get("limit_key") + input_token = kwargs.get(next_token_name) + limit = kwargs.get(limit_name, None) + # Remove pagination arguments from our input kwargs + # We need this to verify that our input kwargs are the same across invocations + # list_all(service="x") next_token = "a" + # list_all(service="x", next_token="a") ==> Works fine + # list_all(service="y", next_token="a") ==> Should throw an error, as the input_kwargs are different + input_kwargs = deepcopy(kwargs) + input_kwargs.pop(next_token_name, None) + input_kwargs.pop(limit_name, None) fail_on_invalid_token = pagination_config.get("fail_on_invalid_token", True) paginator = Paginator( max_results=limit, max_results_default=pagination_config.get("limit_default"), starting_token=input_token, - page_ending_range_keys=pagination_config.get("page_ending_range_keys"), - param_values_to_check=kwargs, + unique_attribute=pagination_config.get("unique_attribute"), + param_values_to_check=input_kwargs, fail_on_invalid_token=fail_on_invalid_token, ) + + # Determine which parameters to pass + (arg_names, _, has_kwargs, _, _, _, _) = inspect.getfullargspec(func) + # If the target-func expects `**kwargs`, we can pass everything + if not has_kwargs: + # If the target-function does not expect the next_token/limit, do not pass it + if next_token_name not in arg_names: + kwargs.pop(next_token_name, None) + if limit_name not in arg_names: + kwargs.pop(limit_name, None) + results = func(*args, **kwargs) return paginator.paginate(results) @@ -47,13 +70,15 @@ class Paginator(object): max_results=None, max_results_default=None, starting_token=None, - page_ending_range_keys=None, + unique_attribute=None, param_values_to_check=None, fail_on_invalid_token=True, ): self._max_results = max_results if max_results else max_results_default self._starting_token = starting_token - self._page_ending_range_keys = page_ending_range_keys + self._unique_attributes = unique_attribute + if not isinstance(unique_attribute, list): + self._unique_attributes = [unique_attribute] self._param_values_to_check = param_values_to_check self._fail_on_invalid_token = fail_on_invalid_token self._token_encoder = TokenEncoder() @@ -69,8 +94,7 @@ class Paginator(object): try: next_token = self._token_decoder.decode(next_token) except (ValueError, TypeError, UnicodeDecodeError): - if self._fail_on_invalid_token: - raise InvalidToken("Invalid token") + self._raise_exception_if_required(next_token) return None if next_token.get("parameterChecksum") != self._param_checksum: raise InvalidToken( @@ -78,20 +102,40 @@ class Paginator(object): ) return next_token + def _raise_exception_if_required(self, token): + if self._fail_on_invalid_token: + if isinstance(self._fail_on_invalid_token, type): + # we need to raise a custom exception + func_info = inspect.getfullargspec(self._fail_on_invalid_token) + arg_names, _, _, _, kwarg_names, _, _ = func_info + # arg_names == [self] or [self, token_argument_that_can_have_any_name] + requires_token_arg = len(arg_names) > 1 + if requires_token_arg: + raise self._fail_on_invalid_token(token) + else: + raise self._fail_on_invalid_token() + raise InvalidToken("Invalid token") + def _calculate_parameter_checksum(self): - if not self._param_values_to_check: - return None - return reduce( - lambda x, y: x ^ y, - [hash(item) for item in self._param_values_to_check.items()], - ) + def freeze(o): + if not o: + return None + if isinstance(o, dict): + return frozenset({k: freeze(v) for k, v in o.items()}.items()) + + if isinstance(o, (list, tuple, set)): + return tuple([freeze(v) for v in o]) + + return o + + return hash(freeze(self._param_values_to_check)) def _check_predicate(self, item): if self._parsed_token is None: return False - page_ending_range_key = self._parsed_token["pageEndingRangeKey"] - predicate_values = page_ending_range_key.split("|") - for (index, attr) in enumerate(self._page_ending_range_keys): + unique_attributes = self._parsed_token["uniqueAttributes"] + predicate_values = unique_attributes.split("|") + for (index, attr) in enumerate(self._unique_attributes): curr_val = item[attr] if type(item) == dict else getattr(item, attr, None) if not curr_val == predicate_values[index]: return False @@ -102,12 +146,12 @@ class Paginator(object): if self._param_checksum: token_dict["parameterChecksum"] = self._param_checksum range_keys = [] - for (index, attr) in enumerate(self._page_ending_range_keys): + for (index, attr) in enumerate(self._unique_attributes): if type(next_item) == dict: range_keys.append(next_item[attr]) else: range_keys.append(getattr(next_item, attr)) - token_dict["pageEndingRangeKey"] = "|".join(range_keys) + token_dict["uniqueAttributes"] = "|".join(range_keys) return self._token_encoder.encode(token_dict) def paginate(self, results): @@ -133,6 +177,6 @@ class Paginator(object): next_token = None if results_page and index_end < len(results): - page_ending_result = results[index_end] - next_token = self._build_next_token(page_ending_result) + last_resource_on_this_page = results[index_end] + next_token = self._build_next_token(last_resource_on_this_page) return results_page, next_token diff --git a/tests/test_utilities/test_paginator.py b/tests/test_utilities/test_paginator.py new file mode 100644 index 000000000..f12ff7f08 --- /dev/null +++ b/tests/test_utilities/test_paginator.py @@ -0,0 +1,299 @@ +import unittest + +import pytest +import sure # noqa # pylint: disable=unused-import +from moto.utilities.paginator import Paginator, paginate +from moto.core.exceptions import InvalidToken + + +results = [ + {"id": f"id{i}", "name": f"name{i}", "arn": f"arn:aws:thing/name{i}"} + for i in range(0, 10) +] + + +class Model: + def __init__(self, i): + self.id = f"id{i}" + self.name = f"name{i}" + self.arn = f"arn:aws:thing/{self.name}" + + +model_results = [Model(i) for i in range(0, 100)] + + +def test_paginator_without_max_results__throws_error(): + p = Paginator() + with pytest.raises(TypeError): + p.paginate(results) + + +def test_paginator__paginate_with_just_max_results(): + p = Paginator(max_results=50) + resp = p.paginate(results) + resp.should.have.length_of(2) + + page, next_token = resp + next_token.should.equal(None) + page.should.equal(results) + + +def test_paginator__paginate_without_range_key__throws_error(): + p = Paginator(max_results=2) + with pytest.raises(KeyError): + p.paginate(results) + + +def test_paginator__paginate_with_unknown_range_key__throws_error(): + p = Paginator(max_results=2, unique_attribute=["unknown"]) + with pytest.raises(KeyError): + p.paginate(results) + + +def test_paginator__paginate_5(): + p = Paginator(max_results=5, unique_attribute=["name"]) + resp = p.paginate(results) + resp.should.have.length_of(2) + + page, next_token = resp + next_token.shouldnt.equal(None) + page.should.equal(results[0:5]) + + +def test_paginator__paginate_5__use_different_range_keys(): + p = Paginator(max_results=5, unique_attribute="name") + _, token_as_str = p.paginate(results) + + p = Paginator(max_results=5, unique_attribute=["name"]) + _, token_as_lst = p.paginate(results) + + token_as_lst.shouldnt.be(None) + token_as_lst.should.equal(token_as_str) + + p = Paginator(max_results=5, unique_attribute=["name", "arn"]) + _, token_multiple = p.paginate(results) + token_multiple.shouldnt.be(None) + token_multiple.shouldnt.equal(token_as_str) + + +def test_paginator__paginate_twice(): + p = Paginator(max_results=5, unique_attribute=["name"]) + resp = p.paginate(results) + resp.should.have.length_of(2) + + page, next_token = resp + + p = Paginator(max_results=10, unique_attribute=["name"], starting_token=next_token) + resp = p.paginate(results) + + page, next_token = resp + next_token.should.equal(None) + page.should.equal(results[5:]) + + +def test_paginator__invalid_token(): + with pytest.raises(InvalidToken): + Paginator(max_results=5, unique_attribute=["name"], starting_token="unknown") + + +def test_paginator__invalid_token__but_we_just_dont_care(): + p = Paginator( + max_results=5, + unique_attribute=["name"], + starting_token="unknown", + fail_on_invalid_token=False, + ) + res, token = p.paginate(results) + + res.should.equal([]) + token.should.equal(None) + + +class CustomInvalidTokenException(BaseException): + def __init__(self, token): + self.message = f"Invalid token: {token}" + + +class GenericInvalidTokenException(BaseException): + def __init__(self): + self.message = "Invalid token!" + + +class TestDecorator(unittest.TestCase): + PAGINATION_MODEL = { + "method_returning_dict": { + "input_token": "next_token", + "limit_key": "max_results", + "limit_default": 100, + "unique_attribute": "name", + }, + "method_returning_instances": { + "input_token": "next_token", + "limit_key": "max_results", + "limit_default": 10, + "limit_max": 50, + "unique_attribute": "name", + }, + "method_returning_args": { + "limit_key": "max_results", + "unique_attribute": "name", + }, + "method_specifying_invalidtoken_exception": { + "limit_key": "max_results", + "limit_default": 5, + "unique_attribute": "name", + "fail_on_invalid_token": CustomInvalidTokenException, + }, + "method_specifying_generic_invalidtoken_exception": { + "limit_key": "max_results", + "limit_default": 5, + "unique_attribute": "name", + "fail_on_invalid_token": GenericInvalidTokenException, + }, + "method_expecting_token_as_kwarg": { + "input_token": "custom_token", + "limit_default": 1, + "unique_attribute": "name", + }, + "method_expecting_limit_as_kwarg": { + "limit_key": "custom_limit", + "limit_default": 1, + "unique_attribute": "name", + }, + "method_with_list_as_kwarg": {"limit_default": 1, "unique_attribute": "name",}, + } + + @paginate(pagination_model=PAGINATION_MODEL) + def method_returning_dict(self): + return results + + @paginate(pagination_model=PAGINATION_MODEL) + def method_returning_instances(self): + return model_results + + @paginate(pagination_model=PAGINATION_MODEL) + def method_without_configuration(self): + return results + + @paginate(pagination_model=PAGINATION_MODEL) + def method_returning_args(self, *args, **kwargs): + return [*args] + [(k, v) for k, v in kwargs.items()] + + @paginate(pagination_model=PAGINATION_MODEL) + def method_expecting_token_as_kwarg(self, custom_token=None): + self.custom_token = custom_token + return [{"name": "item1"}, {"name": "item2"}] + + @paginate(pagination_model=PAGINATION_MODEL) + def method_expecting_limit_as_kwarg(self, custom_limit): + self.custom_limit = custom_limit + return [{"name": "item1"}, {"name": "item2"}] + + @paginate(pagination_model=PAGINATION_MODEL) + def method_with_list_as_kwarg(self, resources=[]): + return resources or results + + @paginate(PAGINATION_MODEL) + def method_specifying_invalidtoken_exception(self): + return results + + @paginate(PAGINATION_MODEL) + def method_specifying_generic_invalidtoken_exception(self): + return results + + def test__method_returning_dict(self): + page, token = self.method_returning_dict() + page.should.equal(results) + token.should.equal(None) + + def test__method_returning_instances(self): + page, token = self.method_returning_instances() + page.should.equal(model_results[0:10]) + token.shouldnt.equal(None) + + def test__method_without_configuration(self): + with pytest.raises(ValueError): + self.method_without_configuration() + + def test__input_arguments_are_returned(self): + resp, token = self.method_returning_args(1, "2", next_token=None, max_results=5) + resp.should.have.length_of(4) + resp.should.contain(1) + resp.should.contain("2") + resp.should.contain(("next_token", None)) + resp.should.contain(("max_results", 5)) + token.should.equal(None) + + def test__pass_exception_on_invalid_token(self): + # works fine if no token is specified + self.method_specifying_invalidtoken_exception() + + # throws exception if next_token is invalid + with pytest.raises(CustomInvalidTokenException) as exc: + self.method_specifying_invalidtoken_exception( + next_token="some invalid token" + ) + exc.value.should.be.a(CustomInvalidTokenException) + exc.value.message.should.equal("Invalid token: some invalid token") + + def test__pass_generic_exception_on_invalid_token(self): + # works fine if no token is specified + self.method_specifying_generic_invalidtoken_exception() + + # throws exception if next_token is invalid + # Exception does not take any arguments - our paginator needs to verify whether the next_token arg is expected + with pytest.raises(GenericInvalidTokenException) as exc: + self.method_specifying_generic_invalidtoken_exception( + next_token="some invalid token" + ) + exc.value.should.be.a(GenericInvalidTokenException) + exc.value.message.should.equal("Invalid token!") + + def test__invoke_function_that_expects_token_as_keyword(self): + resp, first_token = self.method_expecting_token_as_kwarg() + resp.should.equal([{"name": "item1"}]) + first_token.shouldnt.equal(None) + self.custom_token.should.equal(None) + + # Verify the custom_token is received in the business method + # Could be handy for additional validation + resp, token = self.method_expecting_token_as_kwarg(custom_token=first_token) + self.custom_token.should.equal(first_token) + + def test__invoke_function_that_expects_limit_as_keyword(self): + self.method_expecting_limit_as_kwarg(custom_limit=None) + self.custom_limit.should.equal(None) + + # Verify the custom_limit is received in the business method + # Could be handy for additional validation + self.method_expecting_limit_as_kwarg(custom_limit=1) + self.custom_limit.should.equal(1) + + def test__verify_kwargs_can_be_a_list(self): + # Use case - verify that the kwarg can be of type list + # Paginator creates a hash for all kwargs + # We need to be make sure that the hash-function can deal with lists + resp, token = self.method_with_list_as_kwarg() + resp.should.equal(results[0:1]) + + resp, token = self.method_with_list_as_kwarg(next_token=token) + resp.should.equal(results[1:2]) + + custom_list = [{"name": "a"}, {"name": "b"}] + resp, token = self.method_with_list_as_kwarg(resources=custom_list) + resp.should.equal(custom_list[0:1]) + + resp, token = self.method_with_list_as_kwarg( + resources=custom_list, next_token=token + ) + resp.should.equal(custom_list[1:]) + token.should.equal(None) + + def test__paginator_fails_with_inconsistent_arguments(self): + custom_list = [{"name": "a"}, {"name": "b"}] + resp, token = self.method_with_list_as_kwarg(resources=custom_list) + resp.should.equal(custom_list[0:1]) + + with pytest.raises(InvalidToken): + # This should fail, as our 'resources' argument is inconsistent with the original resources that were provided + self.method_with_list_as_kwarg(resources=results, next_token=token)