# Kinesis Aggregation/Deaggregation Libraries for Python
#
# Copyright 2014, Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Amazon Software License (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#
# http://aws.amazon.com/asl/
#
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.

from __future__ import print_function
import aws_kinesis_agg.deaggregator as deagg
import base64
import json
import unittest

# This is an aggregated record generated by the actual KPL library (with explicitly specified EHK values)
# (see com.amazonaws.kinesis.producer.SampleKPLProducer in the Java part of this project)
kpl_generated_rec_with_ehks = {
    "kinesis": {
        "kinesisSchemaVersion": "1.0",
        "partitionKey": "a",
        "sequenceNumber": "49581544954143914480256700934210366840104395343224897586",
        "data": "84mawgokZmMwM2RkODgtM2U3OS00NDhhLWIwMWEtN2NmMWJkNDdiNzg0CiRjYWU0MWIxYy1lYTYxLTQzZjItOTBiZS1iODc1NWViZj"
                "g4ZTIKJGQ0OTA2OTBjLWU3NGQtNGRiMi1hM2M4LWQ4ZjJmMTg0ZmQyMwokYzkyNGJjMDktYjg1ZS00N2YxLWIzMmUtMzM2NTIyZWU1"
                "M2M4EiYzODQ4NjQ5NTg2NzUwODM5OTA3ODE1OTcyMzg0NjA1MTgwNzAyMBInMTkzNzg3NjAwMDM3NjgxNzA2OTUyMTQzMzU3MDcxOT"
                "E2MzUyNjA0EicyNjY4ODA0MzY5NjQ5MzI0MjQyNjU0NjY5MTY3MzQwNjg2ODQ0MzkSJzMzOTYwNjYwMDk0Mjk2NzM5MTg1NDYwMzU1"
                "MjQwMjAyMTg0NzI5MhomCAAQABogUkVDT1JEIDIyIHBlZW9iaGN6YnpkbXNrYm91cGd5cQoaJggBEAEaIFJFQ09SRCAyMyB1c3dreG"
                "Z0eHJvZXVzc2N4c2pobm8KGiYIAhACGiBSRUNPUkQgMjQgY2FzZWhkZ2l2ZmF4ZXVzdGx5c3p5ChomCAMQAxogUkVDT1JEIDI1IG52"
                "ZmZ2cG11b2dkb3BqaGFtZXZyawpRwVPQ3go0yp4Y6kvM0q3V",
        "approximateArrivalTimestamp": 1518133507.063
    },
    "eventSource": "aws:kinesis",
    "eventVersion": "1.0",
    "eventID": "shardId-000000000003:49581544954143914480256700934210366840104395343224897586",
    "eventName": "aws:kinesis:record",
    "invokeIdentityArn": "arn:aws:iam::857565855790:role/service-role/Python27KinesisDeaggregatorRole",
    "awsRegion": "us-east-1",
    "eventSourceARN": "arn:aws:kinesis:us-east-1:857565855790:stream/AggRecordStream"
}

# This is an aggregated record generated by the actual KPL library (with EHK set to null)
# (see com.amazonaws.kinesis.producer.SampleKPLProducer in the Java part of this project)
kpl_generated_rec_no_ehks = {
    'kinesis': {
        'eventVersion': '1.0',
        'eventID': 'shardId-000000000003:49581544954143914480259384464784359468274789083712061490',
        'approximateArrivalTimestamp': 1520372221.837,
        'partitionKey': 'a',
        'data': '84mawgokNzhmZWIxMmEtZGJhMC00NWNhLWE5MWUtYmIxZjJmMTgxOWI0CiRjMmE3NTc4Ny00NjliLTQzMTAtODQwZC1kNDg2ZGNhN2'
                'ViNWUKJDY1MGU5MzYyLTU3MmItNDQyNy1iM2ZjLTEzNTQ5ZDdlNWFlNQokMWQ4ZDk2MDAtMDBiNy00NmYzLWE5ODMtZGU4MzQ3NzU0'
                'MGMwCiRmMmU3MThhMC0zODliLTQ5NGEtYjc5Ni0zMzU1YjA3NTY5Y2MKJDk2NWI4ZGE0LWE2YmQtNDc1NS04MWM5LWU3MTgxYWI3ZG'
                'M5YQokZDY5M2M1ZjAtNTc3Mi00NmM5LThkODYtMDhjNzA1NWJkYjc1CiQ4ZDkyNDc2Yi1lYjg5LTRlODEtOTlmYi1jMmJhNjNhZDZk'
                'OTAaJAgAGiBSRUNPUkQgMTggaWtscnBzdnloeG5lcm9kZmVoZXJ1ChokCAEaIFJFQ09SRCAxOSBsbm1udWR5aHZwdmNncm9rZm16ZG'
                'gKGiQIAhogUkVDT1JEIDIwIHBqcXJjZmJ4YWZ6ZG5kem1iZ2FuZQoaJAgDGiBSRUNPUkQgMjEgYWFmcmV5cm50d3V3eGxpeWdrYmhr'
                'ChokCAQaIFJFQ09SRCAyMiBxc2FhZGhic2dtdnFnZGlmZ3V5b3UKGiQIBRogUkVDT1JEIDIzIHdqb2tvaXd2enpobGZ0bHpocWZqaw'
                'oaJAgGGiBSRUNPUkQgMjQgem9wY2FzY3J0aGJicG5qaXd1aGpiChokCAcaIFJFQ09SRCAyNSB0eGJidm5venJ3c2JveHVvbXFib3EK'
                'vV1rrLdU+Sy2v8xkgZ5YaA==',
        'kinesisSchemaVersion': '1.0',
        'sequenceNumber': '49581544954143914480259384464784359468274789083712061490'
    },
    'invokeIdentityArn': 'arn:aws:iam::857565855790:role/service-role/Python27KinesisDeaggregatorRole',
    'eventName': 'aws:kinesis:record',
    'eventSourceARN': 'arn:aws:kinesis:us-east-1:857565855790:stream/AggRecordStream',
    'eventSource': 'aws:kinesis',
    'awsRegion': 'us-east-1'
}

# This is an aggregated record generated by the actual KPL library but then passed through Kinesis Analytics and the
# record is as it's seen by the Kinesis Analytics preprocessor lambda function.
kpl_generated_ka_rec = {
   'recordId': '49597411459012285111935017621194032819223185658889633794',
   'kinesisStreamRecordMetadata': {
      'approximateArrivalTimestamp': 1562602080951,
      'sequenceNumber': '49597411459012285111935017621194032819223185658889633794',
      'partitionKey': 'a',
      'shardId': 'shardId-000000000000'
   },
   'data': '84mawgoNMTU2MjYwMjA3NDg5NhInMzM5NjA2NjAwOTQyOTY3MzkxODU0NjAzNTUyNDAyMDIxODQ3MjkyGicIABAAGiFSRUNPUkQgMjAwNSB'
           'wanFyY2ZieGFmemRuZHptYmdhbmUaJwgAEAAaIVJFQ09SRCAyMDA2IHBqcXJjZmJ4YWZ6ZG5kem1iZ2FuZRonCAAQABohUkVDT1JEIDIwMD'
           'cgcGpxcmNmYnhhZnpkbmR6bWJnYW5lGicIABAAGiFSRUNPUkQgMjAwOCBwanFyY2ZieGFmemRuZHptYmdhbmUHrLqXAgpJOT8v1lp81gM/'
}

# This is an non-aggregated record generated by the actual KPL library but then passed through Kinesis Analytics and the
# record is as it's seen by the Kinesis Analytics preprocessor lambda function.  Kinesis Analytics will mix aggregated
# and non aggregated records together and they need to be handled properly.
kpl_generated_ka_non_agg_rec = {
   'recordId': '49597411459012285111935017620416693517210978893395132418',
   'kinesisStreamRecordMetadata': {
      'approximateArrivalTimestamp': 1562602078827,
      'sequenceNumber': '49597411459012285111935017620416693517210978893395132418',
      'partitionKey': '1562602074896',
      'shardId': 'shardId-000000000000'
   },
   'data': 'UkVDT1JEIDc0OSBwanFyY2ZieGFmemRuZHptYmdhbmU='
}


class RecordDeaggregatorTest(unittest.TestCase):

    def test_deagg_with_real_kpl_data_with_ehks(self):

        # Each entry is of the form (pk, ehk, data) and corresponds to the expected results
        actual_user_records = [
            {
                'partitionKey': 'fc03dd88-3e79-448a-b01a-7cf1bd47b784',
                'explicitHashKey': '38486495867508399078159723846051807020',
                'data': 'RECORD 22 peeobhczbzdmskboupgyq\n'
            },
            {
                'partitionKey': 'cae41b1c-ea61-43f2-90be-b8755ebf88e2',
                'explicitHashKey': '193787600037681706952143357071916352604',
                'data': 'RECORD 23 uswkxftxroeusscxsjhno\n'
            },
            {
                'partitionKey': 'd490690c-e74d-4db2-a3c8-d8f2f184fd23',
                'explicitHashKey': '266880436964932424265466916734068684439',
                'data': 'RECORD 24 casehdgivfaxeustlyszy\n'
            },
            {
                'partitionKey': 'c924bc09-b85e-47f1-b32e-336522ee53c8',
                'explicitHashKey': '339606600942967391854603552402021847292',
                'data': 'RECORD 25 nvffvpmuogdopjhamevrk\n'
            }
        ]

        records = deagg.deaggregate_records(json.loads(json.dumps(kpl_generated_rec_with_ehks)))

        self.assertEqual(4, len(records), 'Deaggregated the wrong number of records aggregated record.')

        for i in range(0, len(records)):
            generated_record = records[i]
            actual_record = actual_user_records[i]

            self.assertEqual(generated_record['kinesis']['partitionKey'], actual_record['partitionKey'],
                             'Actual and generated partition keys do not match for record %d' % i)
            self.assertEqual(generated_record['kinesis']['explicitHashKey'], actual_record['explicitHashKey'],
                             'Actual and generated explicit hash keys do not match for record %d' % i)

            # Decode base64 to bytes and decode bytes to utf-8
            decoded_data = base64.b64decode(generated_record['kinesis']['data']).decode('utf-8')
            actual_data = actual_record['data']
            self.assertEqual(decoded_data, actual_data,
                             'Deaggregated data does not match expected actual data.')

    def test_deagg_with_real_kpl_data_no_ehks(self):

        # Each entry is of the form (pk, ehk, data) and corresponds to the expected results
        actual_user_records = [
            {
                'partitionKey': '78feb12a-dba0-45ca-a91e-bb1f2f1819b4',
                'data': 'RECORD 18 iklrpsvyhxnerodfeheru\n'
            },
            {
                'partitionKey': 'c2a75787-469b-4310-840d-d486dca7eb5e',
                'data': 'RECORD 19 lnmnudyhvpvcgrokfmzdh\n'
            },
            {
                'partitionKey': '650e9362-572b-4427-b3fc-13549d7e5ae5',
                'data': 'RECORD 20 pjqrcfbxafzdndzmbgane\n'
            },
            {
                'partitionKey': '1d8d9600-00b7-46f3-a983-de83477540c0',
                'data': 'RECORD 21 aafreyrntwuwxliygkbhk\n'
            },
            {
                'partitionKey': 'f2e718a0-389b-494a-b796-3355b07569cc',
                'data': 'RECORD 22 qsaadhbsgmvqgdifguyou\n'
            },
            {
                'partitionKey': '965b8da4-a6bd-4755-81c9-e7181ab7dc9a',
                'data': 'RECORD 23 wjokoiwvzzhlftlzhqfjk\n'
            },
            {
                'partitionKey': 'd693c5f0-5772-46c9-8d86-08c7055bdb75',
                'data': 'RECORD 24 zopcascrthbbpnjiwuhjb\n'
            },
            {
                'partitionKey': '8d92476b-eb89-4e81-99fb-c2ba63ad6d90',
                'data': 'RECORD 25 txbbvnozrwsboxuomqboq\n'
            }
        ]

        records = deagg.deaggregate_records(json.loads(json.dumps(kpl_generated_rec_no_ehks)))

        self.assertEqual(8, len(records), 'Deaggregated the wrong number of records aggregated record.')

        for i in range(0, len(records)):
            generated_record = records[i]
            actual_record = actual_user_records[i]

            self.assertEqual(generated_record['kinesis']['partitionKey'], actual_record['partitionKey'],
                             'Actual and generated partition keys do not match for record %d.' % i)
            self.assertIsNone(generated_record['kinesis']['explicitHashKey'])

            decoded_data = base64.b64decode(generated_record['kinesis']['data']).decode('utf-8')
            actual_data = actual_record['data']

            self.assertEqual(decoded_data, actual_data,
                             'Deaggregated data does not match expected actual data for record %d.' % i)

    def test_deagg_with_real_kpl_generated_ka_rec_with_ehks(self):

        # Each entry is of the form (pk, ehk, data) and corresponds to the expected results
        actual_user_records = [
           {
             'explicitHashKey': '339606600942967391854603552402021847292',
             'partitionKey': '1562602074896',
             'data': 'RECORD 2005 pjqrcfbxafzdndzmbgane'
           },
           {
             'explicitHashKey': '339606600942967391854603552402021847292',
             'partitionKey': '1562602074896',
             'data': 'RECORD 2006 pjqrcfbxafzdndzmbgane'
           },
           {
             'explicitHashKey': '339606600942967391854603552402021847292',
             'partitionKey': '1562602074896',
             'data': 'RECORD 2007 pjqrcfbxafzdndzmbgane'
           },
           {
             'explicitHashKey': '339606600942967391854603552402021847292',
             'partitionKey': '1562602074896',
             'data': 'RECORD 2008 pjqrcfbxafzdndzmbgane'
           }
        ]

        records = deagg.deaggregate_records(json.loads(json.dumps(kpl_generated_ka_rec)))

        self.assertEqual(4, len(records), 'Deaggregated the wrong number of records aggregated record.')

        stuff = []
        for i in range(0, len(records)):
            generated_record = records[i]
            actual_record = actual_user_records[i]

            stuff.append(generated_record)

            self.assertEqual(generated_record['kinesis']['partitionKey'], actual_record['partitionKey'],
                             'Actual and generated partition keys do not match for record %d.' % i)

            self.assertEqual(generated_record['kinesis']['explicitHashKey'], actual_record['explicitHashKey'],
                             'Actual and generated explicit hash keys do not match for record %d' % i)

            decoded_data = base64.b64decode(generated_record['kinesis']['data']).decode('utf-8')
            actual_data = actual_record['data']

            self.assertEqual(decoded_data, actual_data,
                             'Deaggregated data does not match expected actual data for record %d.' % i)

    def test_deagg_with_real_kpl_generated_ka_non_agg_rec_no_ehks(self):

        # Each entry is of the form (pk, ehk, data) and corresponds to the expected results
        actual_user_records = [
           {
             'partitionKey': '1562602074896',
             'data': 'RECORD 749 pjqrcfbxafzdndzmbgane'
           }
        ]

        records = deagg.deaggregate_records(json.loads(json.dumps(kpl_generated_ka_non_agg_rec)))

        self.assertEqual(1, len(records), 'Deaggregated the wrong number of records aggregated record.')

        stuff = []
        for i in range(0, len(records)):
            generated_record = records[i]
            actual_record = actual_user_records[i]

            stuff.append(generated_record)

            self.assertEqual(generated_record['kinesis']['partitionKey'], actual_record['partitionKey'],
                             'Actual and generated partition keys do not match for record %d.' % i)

            decoded_data = base64.b64decode(generated_record['kinesis']['data']).decode('utf-8')
            actual_data = actual_record['data']

            self.assertEqual(decoded_data, actual_data,
                             'Deaggregated data does not match expected actual data for record %d.' % i)


if __name__ == '__main__':
    unittest.main()
