diff --git a/moto/kinesis/utils.py b/moto/kinesis/utils.py index 337728f02..0c3edbb5a 100644 --- a/moto/kinesis/utils.py +++ b/moto/kinesis/utils.py @@ -1,8 +1,19 @@ +import sys import base64 from .exceptions import InvalidArgumentError +if sys.version_info[0] == 2: + encode_method = base64.encodestring + decode_method = base64.decodestring +elif sys.version_info[0] == 3: + encode_method = base64.encodebytes + decode_method = base64.decodebytes +else: + raise Exception("Python version is not supported") + + def compose_new_shard_iterator(stream_name, shard, shard_iterator_type, starting_sequence_number, at_timestamp): if shard_iterator_type == "AT_SEQUENCE_NUMBER": @@ -22,7 +33,7 @@ def compose_new_shard_iterator(stream_name, shard, shard_iterator_type, starting def compose_shard_iterator(stream_name, shard, last_sequence_id): - return base64.encodestring( + return encode_method( "{0}:{1}:{2}".format( stream_name, shard.shard_id, @@ -32,4 +43,4 @@ def compose_shard_iterator(stream_name, shard, last_sequence_id): def decompose_shard_iterator(shard_iterator): - return base64.decodestring(shard_iterator.encode("utf-8")).decode("utf-8").split(":") + return decode_method(shard_iterator.encode("utf-8")).decode("utf-8").split(":")