def test_stream_consumers(self): client = aws_stack.connect_to_service('kinesis') stream_name = 'test-%s' % short_uid() stream_arn = aws_stack.kinesis_stream_arn(stream_name) def assert_consumers(count): consumers = client.list_stream_consumers( StreamARN=stream_arn).get('Consumers') self.assertEqual(len(consumers), count) return consumers # create stream and assert 0 consumers client.create_stream(StreamName=stream_name, ShardCount=1) assert_consumers(0) # create consumer and assert 1 consumer consumer_name = 'cons1' client.register_stream_consumer(StreamARN=stream_arn, ConsumerName=consumer_name) consumers = assert_consumers(1) self.assertEqual(consumers[0]['ConsumerName'], consumer_name) self.assertIn('/%s' % consumer_name, consumers[0]['ConsumerARN']) # delete non-existing consumer and assert 1 consumer client.deregister_stream_consumer(StreamARN=stream_arn, ConsumerName='_invalid_') assert_consumers(1) # delete existing consumer and assert 0 remaining consumers client.deregister_stream_consumer(StreamARN=stream_arn, ConsumerName=consumer_name) assert_consumers(0)
def find_stream_for_consumer(consumer_arn): kinesis = aws_stack.connect_to_service('kinesis') for stream_name in kinesis.list_streams()['StreamNames']: stream_arn = aws_stack.kinesis_stream_arn(stream_name) for cons in kinesis.list_stream_consumers(StreamARN=stream_arn)['Consumers']: if cons['ConsumerARN'] == consumer_arn: return stream_name raise Exception('Unable to find stream for stream consumer %s' % consumer_arn)
def test_subscribe_to_shard_with_sequence_number_as_iterator(self): client = aws_stack.create_external_boto_client("kinesis") stream_name = "test-%s" % short_uid() stream_arn = aws_stack.kinesis_stream_arn(stream_name) record_data = "Hello world" # create stream and consumer result = client.create_stream(StreamName=stream_name, ShardCount=1) sleep(1) result = client.register_stream_consumer(StreamARN=stream_arn, ConsumerName="c1")["Consumer"] sleep(1) # get starting sequence number response = client.describe_stream(StreamName=stream_name) sequence_number = ( response.get("StreamDescription").get("Shards")[0].get( "SequenceNumberRange").get("StartingSequenceNumber")) # subscribe to shard with iterator type as AT_SEQUENCE_NUMBER response = client.describe_stream(StreamName=stream_name) shard_id = response.get("StreamDescription").get("Shards")[0].get( "ShardId") result = client.subscribe_to_shard( ConsumerARN=result["ConsumerARN"], ShardId=shard_id, StartingPosition={ "Type": "AT_SEQUENCE_NUMBER", "SequenceNumber": sequence_number, }, ) stream = result["EventStream"] # put records num_records = 5 for i in range(num_records): client.put_records( StreamName=stream_name, Records=[{ "Data": record_data, "PartitionKey": "1" }], ) results = [] for entry in stream: records = entry["SubscribeToShardEvent"]["Records"] results.extend(records) if len(results) >= num_records: break # assert results self.assertEqual(num_records, len(results)) for record in results: self.assertEqual(str.encode(record_data), record["Data"]) # clean up client.deregister_stream_consumer(StreamARN=stream_arn, ConsumerName="c1") client.delete_stream(StreamName=stream_name, EnforceConsumerDeletion=True)
def test_subscribe_to_shard(self): client = aws_stack.connect_to_service("kinesis") stream_name = "test-%s" % short_uid() stream_arn = aws_stack.kinesis_stream_arn(stream_name) # create stream and consumer result = client.create_stream(StreamName=stream_name, ShardCount=1) sleep(1) result = client.register_stream_consumer(StreamARN=stream_arn, ConsumerName="c1")["Consumer"] sleep(1) # subscribe to shard response = client.describe_stream(StreamName=stream_name) shard_id = response.get("StreamDescription").get("Shards")[0].get( "ShardId") result = client.subscribe_to_shard( ConsumerARN=result["ConsumerARN"], ShardId=shard_id, StartingPosition={"Type": "TRIM_HORIZON"}, ) stream = result["EventStream"] # put records num_records = 5 msg = b"Hello world" msg_b64 = base64.b64encode(msg) for i in range(num_records): client.put_records(StreamName=stream_name, Records=[{ "Data": msg_b64, "PartitionKey": "1" }]) # assert results results = [] for entry in stream: records = entry["SubscribeToShardEvent"]["Records"] continuation_sequence_number = entry["SubscribeToShardEvent"][ "ContinuationSequenceNumber"] # https://docs.aws.amazon.com/kinesis/latest/APIReference/API_SubscribeToShardEvent.html self.assertIsNotNone( re.fullmatch("^0|([1-9][0-9]{0,128})$", continuation_sequence_number)) results.extend(records) if len(results) >= num_records: break # assert results self.assertEqual(num_records, len(results)) for record in results: self.assertEqual(msg, record["Data"]) # clean up client.deregister_stream_consumer(StreamARN=stream_arn, ConsumerName="c1") client.delete_stream(StreamName=stream_name, EnforceConsumerDeletion=True)
def test_subscribe_to_shard(self): client = aws_stack.connect_to_service('kinesis') stream_name = 'test-%s' % short_uid() stream_arn = aws_stack.kinesis_stream_arn(stream_name) # create stream and consumer result = client.create_stream(StreamName=stream_name, ShardCount=1) sleep(1) result = client.register_stream_consumer(StreamARN=stream_arn, ConsumerName='c1')['Consumer'] sleep(1) # subscribe to shard response = client.describe_stream(StreamName=stream_name) shard_id = response.get('StreamDescription').get('Shards')[0].get( 'ShardId') result = client.subscribe_to_shard( ConsumerARN=result['ConsumerARN'], ShardId=shard_id, StartingPosition={'Type': 'TRIM_HORIZON'}) stream = result['EventStream'] # put records num_records = 5 msg = b'Hello world' msg_b64 = base64.b64encode(msg) for i in range(num_records): client.put_records(StreamName=stream_name, Records=[{ 'Data': msg_b64, 'PartitionKey': '1' }]) # assert results results = [] for entry in stream: records = entry['SubscribeToShardEvent']['Records'] continuation_sequence_number = entry['SubscribeToShardEvent'][ 'ContinuationSequenceNumber'] # https://docs.aws.amazon.com/kinesis/latest/APIReference/API_SubscribeToShardEvent.html self.assertIsNotNone( re.fullmatch('^0|([1-9][0-9]{0,128})$', continuation_sequence_number)) results.extend(records) if len(results) >= num_records: break # assert results self.assertEqual(num_records, len(results)) for record in results: self.assertEqual(msg, record['Data']) # clean up client.deregister_stream_consumer(StreamARN=stream_arn, ConsumerName='c1') client.delete_stream(StreamName=stream_name, EnforceConsumerDeletion=True)
def find_stream_for_consumer(consumer_arn): kinesis = aws_stack.connect_to_service("kinesis") for stream_name in kinesis.list_streams()["StreamNames"]: stream_arn = aws_stack.kinesis_stream_arn(stream_name) for cons in kinesis.list_stream_consumers( StreamARN=stream_arn)["Consumers"]: if cons["ConsumerARN"] == consumer_arn: return stream_name raise Exception("Unable to find stream for stream consumer %s" % consumer_arn)
def test_subscribe_to_shard_with_sequence_number_as_iterator(self): client = aws_stack.connect_to_service('kinesis') stream_name = 'test-%s' % short_uid() stream_arn = aws_stack.kinesis_stream_arn(stream_name) # create stream and consumer result = client.create_stream(StreamName=stream_name, ShardCount=1) sleep(1) result = client.register_stream_consumer(StreamARN=stream_arn, ConsumerName='c1')['Consumer'] sleep(1) # get starting sequence number response = client.describe_stream(StreamName=stream_name) sequence_number = response.get('StreamDescription').get('Shards')[0].get('SequenceNumberRange'). \ get('StartingSequenceNumber') # subscribe to shard with iterator type as AT_SEQUENCE_NUMBER response = client.describe_stream(StreamName=stream_name) shard_id = response.get('StreamDescription').get('Shards')[0].get( 'ShardId') result = client.subscribe_to_shard(ConsumerARN=result['ConsumerARN'], ShardId=shard_id, StartingPosition={ 'Type': 'AT_SEQUENCE_NUMBER', 'SequenceNumber': sequence_number }) stream = result['EventStream'] # put records num_records = 5 for i in range(num_records): client.put_records(StreamName=stream_name, Records=[{ 'Data': 'SGVsbG8gd29ybGQ=', 'PartitionKey': '1' }]) results = [] for entry in stream: records = entry['SubscribeToShardEvent']['Records'] results.extend(records) if len(results) >= num_records: break # assert results self.assertEqual(num_records, len(results)) for record in results: self.assertEqual(b'Hello world', record['Data']) # clean up client.deregister_stream_consumer(StreamARN=stream_arn, ConsumerName='c1') client.delete_stream(StreamName=stream_name, EnforceConsumerDeletion=True)
def test_create_kinesis_event_source_mapping(self): function_name = f"lambda_func-{short_uid()}" stream_name = f"test-foobar-{short_uid()}" testutil.create_lambda_function( handler_file=TEST_LAMBDA_PYTHON_ECHO, func_name=function_name, runtime=LAMBDA_RUNTIME_PYTHON36, ) arn = aws_stack.kinesis_stream_arn(stream_name, account_id="000000000000") lambda_client = aws_stack.create_external_boto_client("lambda") lambda_client.create_event_source_mapping(EventSourceArn=arn, FunctionName=function_name) def process_records(record): assert record aws_stack.create_kinesis_stream(stream_name, delete=True) kinesis_connector.listen_to_kinesis( stream_name=stream_name, listener_func=process_records, wait_until_started=True, ) kinesis = aws_stack.create_external_boto_client("kinesis") stream_summary = kinesis.describe_stream_summary( StreamName=stream_name) self.assertEqual( 1, stream_summary["StreamDescriptionSummary"]["OpenShardCount"]) num_events_kinesis = 10 kinesis.put_records( Records=[{ "Data": "{}", "PartitionKey": "test_%s" % i } for i in range(0, num_events_kinesis)], StreamName=stream_name, ) events = get_lambda_log_events(function_name) self.assertEqual(10, len(events[0]["Records"])) self.assertIn("eventID", events[0]["Records"][0]) self.assertIn("eventSourceARN", events[0]["Records"][0]) self.assertIn("eventSource", events[0]["Records"][0]) self.assertIn("eventVersion", events[0]["Records"][0]) self.assertIn("eventName", events[0]["Records"][0]) self.assertIn("invokeIdentityArn", events[0]["Records"][0]) self.assertIn("awsRegion", events[0]["Records"][0]) self.assertIn("kinesis", events[0]["Records"][0])
def test_stream_consumers(self): client = aws_stack.connect_to_service('kinesis') stream_name = 'test-%s' % short_uid() stream_arn = aws_stack.kinesis_stream_arn(stream_name) def assert_consumers(count): consumers = client.list_stream_consumers(StreamARN=stream_arn).get('Consumers') self.assertEqual(len(consumers), count) return consumers # create stream and assert 0 consumers client.create_stream(StreamName=stream_name, ShardCount=1) assert_consumers(0) # create consumer and assert 1 consumer consumer_name = 'cons1' response = client.register_stream_consumer(StreamARN=stream_arn, ConsumerName=consumer_name) self.assertEqual(response['Consumer']['ConsumerName'], consumer_name) # boto3 converts the timestamp to datetime self.assertTrue(isinstance(response['Consumer']['ConsumerCreationTimestamp'], datetime)) consumers = assert_consumers(1) consumer_arn = consumers[0]['ConsumerARN'] self.assertEqual(consumers[0]['ConsumerName'], consumer_name) self.assertIn('/%s' % consumer_name, consumer_arn) self.assertTrue(isinstance(consumers[0]['ConsumerCreationTimestamp'], datetime)) # lookup stream consumer by describe calls, assert response consumer_description_by_arn = client.describe_stream_consumer( StreamARN=stream_arn, ConsumerARN=consumer_arn)['ConsumerDescription'] self.assertEqual(consumer_description_by_arn['ConsumerName'], consumer_name) self.assertEqual(consumer_description_by_arn['ConsumerARN'], consumer_arn) self.assertEqual(consumer_description_by_arn['StreamARN'], stream_arn) self.assertEqual(consumer_description_by_arn['ConsumerStatus'], 'ACTIVE') self.assertTrue(isinstance(consumer_description_by_arn['ConsumerCreationTimestamp'], datetime)) consumer_description_by_name = client.describe_stream_consumer( StreamARN=stream_arn, ConsumerName=consumer_name)['ConsumerDescription'] self.assertEqual(consumer_description_by_arn, consumer_description_by_name) # delete non-existing consumer and assert 1 consumer client.deregister_stream_consumer(StreamARN=stream_arn, ConsumerName='_invalid_') assert_consumers(1) # delete existing consumer and assert 0 remaining consumers client.deregister_stream_consumer(StreamARN=stream_arn, ConsumerName=consumer_name) assert_consumers(0) # clean up client.delete_stream(StreamName=stream_name)
def process_kinesis_records(records, stream_name): # feed records into listening lambdas try: stream_arn = aws_stack.kinesis_stream_arn(stream_name) sources = get_event_sources(source_arn=stream_arn) for source in sources: arn = source['FunctionArn'] event = {'Records': []} for rec in records: event['Records'].append({'kinesis': rec}) run_lambda(event=event, context={}, func_arn=arn) except Exception as e: LOG.warning('Unable to run Lambda function on Kinesis records: %s %s' % (e, traceback.format_exc()))
def process_kinesis_records(records, stream_name): # feed records into listening lambdas try: stream_arn = aws_stack.kinesis_stream_arn(stream_name) sources = get_event_sources(source_arn=stream_arn) for source in sources: arn = source['FunctionArn'] lambda_function = lambda_arn_to_function[arn] event = {'Records': []} for rec in records: event['Records'].append({'kinesis': rec}) run_lambda(lambda_function, event=event, context={}, func_arn=arn) except Exception, e: print(traceback.format_exc())
def test_subscribe_to_shard(self): client = aws_stack.connect_to_service('kinesis') stream_name = 'test-%s' % short_uid() stream_arn = aws_stack.kinesis_stream_arn(stream_name) # create stream and consumer result = client.create_stream(StreamName=stream_name, ShardCount=1) sleep(1) result = client.register_stream_consumer(StreamARN=stream_arn, ConsumerName='c1')['Consumer'] # subscribe to shard response = client.describe_stream(StreamName=stream_name) shard_id = response.get('StreamDescription').get('Shards')[0].get( 'ShardId') result = client.subscribe_to_shard( ConsumerARN=result['ConsumerARN'], ShardId=shard_id, StartingPosition={'Type': 'TRIM_HORIZON'}) stream = result['EventStream'] # put records num_records = 5 msg = b'Hello world' msg_b64 = base64.b64encode(msg) for i in range(num_records): client.put_records(StreamName=stream_name, Records=[{ 'Data': msg_b64, 'PartitionKey': '1' }]) # assert results results = [] for entry in stream: records = entry['SubscribeToShardEvent']['Records'] results.extend(records) if len(results) >= num_records: break # assert results self.assertEqual(len(results), num_records) for record in results: self.assertEqual(record['Data'], msg) # clean up client.deregister_stream_consumer(StreamARN=stream_arn, ConsumerName='c1') client.delete_stream(StreamName=stream_name)
def process_kinesis_records(records, stream_name): # feed records into listening lambdas try: stream_arn = aws_stack.kinesis_stream_arn(stream_name) sources = get_event_sources(source_arn=stream_arn) for source in sources: arn = source['FunctionArn'] event = { 'Records': [] } for rec in records: event['Records'].append({ 'eventID': 'shardId-000000000000:{0}'.format(rec['sequenceNumber']), 'eventSourceARN': stream_arn, 'kinesis': rec }) run_lambda(event=event, context={}, func_arn=arn) except Exception as e: LOG.warning('Unable to run Lambda function on Kinesis records: %s %s' % (e, traceback.format_exc()))
def test_firehose_kinesis_to_s3(self): kinesis = aws_stack.connect_to_service('kinesis') s3_resource = aws_stack.connect_to_resource('s3') firehose = aws_stack.connect_to_service('firehose') aws_stack.create_kinesis_stream(TEST_STREAM_NAME, delete=True) s3_prefix = '/testdata' test_data = '{"test": "firehose_data_%s"}' % short_uid() # create Firehose stream stream = firehose.create_delivery_stream( DeliveryStreamType='KinesisStreamAsSource', KinesisStreamSourceConfiguration={ 'RoleARN': aws_stack.iam_resource_arn('firehose'), 'KinesisStreamARN': aws_stack.kinesis_stream_arn(TEST_STREAM_NAME) }, DeliveryStreamName=TEST_FIREHOSE_NAME, S3DestinationConfiguration={ 'RoleARN': aws_stack.iam_resource_arn('firehose'), 'BucketARN': aws_stack.s3_bucket_arn(TEST_BUCKET_NAME), 'Prefix': s3_prefix } ) self.assertTrue(stream) self.assertIn(TEST_FIREHOSE_NAME, firehose.list_delivery_streams()['DeliveryStreamNames']) # create target S3 bucket s3_resource.create_bucket(Bucket=TEST_BUCKET_NAME) # put records kinesis.put_record( Data=to_bytes(test_data), PartitionKey='testId', StreamName=TEST_STREAM_NAME ) time.sleep(3) # check records in target bucket all_objects = testutil.list_all_s3_objects() testutil.assert_objects(json.loads(to_str(test_data)), all_objects)
def test_firehose_kinesis_to_s3(self): kinesis = aws_stack.create_external_boto_client("kinesis") s3_resource = aws_stack.connect_to_resource("s3") firehose = aws_stack.create_external_boto_client("firehose") aws_stack.create_kinesis_stream(TEST_STREAM_NAME, delete=True) s3_prefix = "/testdata" test_data = '{"test": "firehose_data_%s"}' % short_uid() # create Firehose stream stream = firehose.create_delivery_stream( DeliveryStreamType="KinesisStreamAsSource", KinesisStreamSourceConfiguration={ "RoleARN": aws_stack.iam_resource_arn("firehose"), "KinesisStreamARN": aws_stack.kinesis_stream_arn(TEST_STREAM_NAME), }, DeliveryStreamName=TEST_FIREHOSE_NAME, S3DestinationConfiguration={ "RoleARN": aws_stack.iam_resource_arn("firehose"), "BucketARN": aws_stack.s3_bucket_arn(TEST_BUCKET_NAME), "Prefix": s3_prefix, }, ) self.assertTrue(stream) self.assertIn(TEST_FIREHOSE_NAME, firehose.list_delivery_streams()["DeliveryStreamNames"]) # create target S3 bucket s3_resource.create_bucket(Bucket=TEST_BUCKET_NAME) # put records kinesis.put_record(Data=to_bytes(test_data), PartitionKey="testId", StreamName=TEST_STREAM_NAME) time.sleep(3) # check records in target bucket all_objects = testutil.list_all_s3_objects() testutil.assert_objects(json.loads(to_str(test_data)), all_objects)
def test_put_events_with_target_kinesis(self, events_client, kinesis_client): rule_name = "rule-{}".format(short_uid()) target_id = "target-{}".format(short_uid()) bus_name = "bus-{}".format(short_uid()) stream_name = "stream-{}".format(short_uid()) stream_arn = aws_stack.kinesis_stream_arn(stream_name) kinesis_client.create_stream(StreamName=stream_name, ShardCount=1) events_client.create_event_bus(Name=bus_name) events_client.put_rule( Name=rule_name, EventBusName=bus_name, EventPattern=json.dumps(TEST_EVENT_PATTERN), ) put_response = events_client.put_targets( Rule=rule_name, EventBusName=bus_name, Targets=[ { "Id": target_id, "Arn": stream_arn, "KinesisParameters": {"PartitionKeyPath": "$.detail-type"}, } ], ) assert "FailedEntryCount" in put_response assert "FailedEntries" in put_response assert put_response["FailedEntryCount"] == 0 assert put_response["FailedEntries"] == [] def check_stream_status(): _stream = kinesis_client.describe_stream(StreamName=stream_name) assert _stream["StreamDescription"]["StreamStatus"] == "ACTIVE" # wait until stream becomes available retry(check_stream_status, retries=7, sleep=0.8) events_client.put_events( Entries=[ { "EventBusName": bus_name, "Source": TEST_EVENT_PATTERN["source"][0], "DetailType": TEST_EVENT_PATTERN["detail-type"][0], "Detail": json.dumps(EVENT_DETAIL), } ] ) stream = kinesis_client.describe_stream(StreamName=stream_name) shard_id = stream["StreamDescription"]["Shards"][0]["ShardId"] shard_iterator = kinesis_client.get_shard_iterator( StreamName=stream_name, ShardId=shard_id, ShardIteratorType="AT_TIMESTAMP", Timestamp=datetime(2020, 1, 1), )["ShardIterator"] record = kinesis_client.get_records(ShardIterator=shard_iterator)["Records"][0] partition_key = record["PartitionKey"] data = json.loads(record["Data"].decode()) assert partition_key == TEST_EVENT_PATTERN["detail-type"][0] assert data["detail"] == EVENT_DETAIL self.assert_valid_event(data)
def get_physical_resource_id(self, attribute=None, **kwargs): return aws_stack.kinesis_stream_arn(self.props.get('Name'))
def test_put_events_with_target_kinesis(self): rule_name = 'rule-{}'.format(short_uid()) target_id = 'target-{}'.format(short_uid()) bus_name = 'bus-{}'.format(short_uid()) stream_name = 'stream-{}'.format(short_uid()) stream_arn = aws_stack.kinesis_stream_arn(stream_name) kinesis_client = aws_stack.connect_to_service('kinesis') kinesis_client.create_stream(StreamName=stream_name, ShardCount=1) self.events_client.create_event_bus(Name=bus_name) self.events_client.put_rule( Name=rule_name, EventBusName=bus_name, EventPattern=json.dumps(TEST_EVENT_PATTERN)) put_response = self.events_client.put_targets( Rule=rule_name, EventBusName=bus_name, Targets=[{ 'Id': target_id, 'Arn': stream_arn, 'KinesisParameters': { 'PartitionKeyPath': '$.detail-type' } }]) self.assertIn('FailedEntryCount', put_response) self.assertIn('FailedEntries', put_response) self.assertEqual(0, put_response['FailedEntryCount']) self.assertEqual([], put_response['FailedEntries']) def put_events(events_client): events_client.put_events( Entries=[{ 'EventBusName': bus_name, 'Source': TEST_EVENT_PATTERN['Source'][0], 'DetailType': TEST_EVENT_PATTERN['detail-type'][0], 'Detail': json.dumps(TEST_EVENT_PATTERN['Detail'][0]) }]) # Stream may be still creating retry(put_events, retries=5, sleep=10, events_client=self.events_client) stream = kinesis_client.describe_stream(StreamName=stream_name) shard_id = stream['StreamDescription']['Shards'][0]['ShardId'] shard_iterator = kinesis_client.get_shard_iterator( StreamName=stream_name, ShardId=shard_id, ShardIteratorType='AT_TIMESTAMP', Timestamp=datetime(2020, 1, 1))['ShardIterator'] record = kinesis_client.get_records( ShardIterator=shard_iterator)['Records'][0] partition_key = record['PartitionKey'] data = json.loads(record['Data'].decode()) self.assertEqual(TEST_EVENT_PATTERN['detail-type'][0], partition_key) self.assertEqual(EVENT_DETAIL, data['detail']) self.assertIsValidEvent(data)
def test_stream_consumers(self): client = aws_stack.create_external_boto_client("kinesis") stream_name = "test-%s" % short_uid() stream_arn = aws_stack.kinesis_stream_arn(stream_name) def assert_consumers(count): consumers = client.list_stream_consumers( StreamARN=stream_arn).get("Consumers") self.assertEqual(count, len(consumers)) return consumers # create stream and assert 0 consumers client.create_stream(StreamName=stream_name, ShardCount=1) sleep(1) assert_consumers(0) # create consumer and assert 1 consumer consumer_name = "cons1" response = client.register_stream_consumer(StreamARN=stream_arn, ConsumerName=consumer_name) sleep(1) self.assertEqual(consumer_name, response["Consumer"]["ConsumerName"]) # boto3 converts the timestamp to datetime self.assertTrue( isinstance(response["Consumer"]["ConsumerCreationTimestamp"], datetime)) consumers = assert_consumers(1) consumer_arn = consumers[0]["ConsumerARN"] self.assertEqual(consumer_name, consumers[0]["ConsumerName"]) self.assertIn("/%s" % consumer_name, consumer_arn) self.assertTrue( isinstance(consumers[0]["ConsumerCreationTimestamp"], datetime)) # lookup stream consumer by describe calls, assert response consumer_description_by_arn = client.describe_stream_consumer( StreamARN=stream_arn, ConsumerARN=consumer_arn)["ConsumerDescription"] self.assertEqual(consumer_name, consumer_description_by_arn["ConsumerName"]) self.assertEqual(consumer_arn, consumer_description_by_arn["ConsumerARN"]) self.assertEqual(stream_arn, consumer_description_by_arn["StreamARN"]) self.assertEqual("ACTIVE", consumer_description_by_arn["ConsumerStatus"]) self.assertTrue( isinstance( consumer_description_by_arn["ConsumerCreationTimestamp"], datetime)) consumer_description_by_name = client.describe_stream_consumer( StreamARN=stream_arn, ConsumerName=consumer_name)["ConsumerDescription"] self.assertEqual(consumer_description_by_arn, consumer_description_by_name) # delete existing consumer and assert 0 remaining consumers client.deregister_stream_consumer(StreamARN=stream_arn, ConsumerName=consumer_name) sleep(1) assert_consumers(0) # clean up client.delete_stream(StreamName=stream_name)
def test_kinesis_lambda_parallelism(self, lambda_client, kinesis_client): function_name = f"lambda_func-{short_uid()}" stream_name = f"test-foobar-{short_uid()}" testutil.create_lambda_function( handler_file=TEST_LAMBDA_PARALLEL_FILE, func_name=function_name, runtime=LAMBDA_RUNTIME_PYTHON36, ) arn = aws_stack.kinesis_stream_arn(stream_name, account_id="000000000000") lambda_client.create_event_source_mapping(EventSourceArn=arn, FunctionName=function_name) def process_records(record): assert record aws_stack.create_kinesis_stream(stream_name, delete=True) kinesis_connector.listen_to_kinesis( stream_name=stream_name, listener_func=process_records, wait_until_started=True, ) kinesis = aws_stack.create_external_boto_client("kinesis") stream_summary = kinesis.describe_stream_summary(StreamName=stream_name) assert 1 == stream_summary["StreamDescriptionSummary"]["OpenShardCount"] num_events_kinesis = 10 # assure async call start = time.perf_counter() kinesis.put_records( Records=[ {"Data": '{"batch": 0}', "PartitionKey": f"test_{i}"} for i in range(0, num_events_kinesis) ], StreamName=stream_name, ) assert (time.perf_counter() - start) < 1 # this should not take more than a second kinesis.put_records( Records=[ {"Data": '{"batch": 1}', "PartitionKey": f"test_{i}"} for i in range(0, num_events_kinesis) ], StreamName=stream_name, ) def get_events(): events = get_lambda_log_events(function_name, regex_filter=r"event.*Records") assert len(events) == 2 return events events = retry(get_events, retries=5) def assertEvent(event, batch_no): assert 10 == len(event["event"]["Records"]) assert "eventID" in event["event"]["Records"][0] assert "eventSourceARN" in event["event"]["Records"][0] assert "eventSource" in event["event"]["Records"][0] assert "eventVersion" in event["event"]["Records"][0] assert "eventName" in event["event"]["Records"][0] assert "invokeIdentityArn" in event["event"]["Records"][0] assert "awsRegion" in event["event"]["Records"][0] assert "kinesis" in event["event"]["Records"][0] assert {"batch": batch_no} == json.loads( base64.b64decode(event["event"]["Records"][0]["kinesis"]["data"]).decode( config.DEFAULT_ENCODING ) ) assertEvent(events[0], 0) assertEvent(events[1], 1) assert (events[1]["executionStart"] - events[0]["executionStart"]) > 5 # cleanup lambda_client.delete_function(FunctionName=function_name) kinesis_client.delete_stream(StreamName=stream_name)
def test_put_events_with_target_kinesis(self): rule_name = "rule-{}".format(short_uid()) target_id = "target-{}".format(short_uid()) bus_name = "bus-{}".format(short_uid()) stream_name = "stream-{}".format(short_uid()) stream_arn = aws_stack.kinesis_stream_arn(stream_name) kinesis_client = aws_stack.connect_to_service("kinesis") kinesis_client.create_stream(StreamName=stream_name, ShardCount=1) self.events_client.create_event_bus(Name=bus_name) self.events_client.put_rule( Name=rule_name, EventBusName=bus_name, EventPattern=json.dumps(TEST_EVENT_PATTERN), ) put_response = self.events_client.put_targets( Rule=rule_name, EventBusName=bus_name, Targets=[{ "Id": target_id, "Arn": stream_arn, "KinesisParameters": { "PartitionKeyPath": "$.detail-type" }, }], ) self.assertIn("FailedEntryCount", put_response) self.assertIn("FailedEntries", put_response) self.assertEqual(0, put_response["FailedEntryCount"]) self.assertEqual([], put_response["FailedEntries"]) def put_events(events_client): events_client.put_events( Entries=[{ "EventBusName": bus_name, "Source": TEST_EVENT_PATTERN["Source"][0], "DetailType": TEST_EVENT_PATTERN["detail-type"][0], "Detail": json.dumps(TEST_EVENT_PATTERN["Detail"][0]), }]) # Stream may be still creating retry(put_events, retries=5, sleep=10, events_client=self.events_client) stream = kinesis_client.describe_stream(StreamName=stream_name) shard_id = stream["StreamDescription"]["Shards"][0]["ShardId"] shard_iterator = kinesis_client.get_shard_iterator( StreamName=stream_name, ShardId=shard_id, ShardIteratorType="AT_TIMESTAMP", Timestamp=datetime(2020, 1, 1), )["ShardIterator"] record = kinesis_client.get_records( ShardIterator=shard_iterator)["Records"][0] partition_key = record["PartitionKey"] data = json.loads(record["Data"].decode()) self.assertEqual(TEST_EVENT_PATTERN["detail-type"][0], partition_key) self.assertEqual(EVENT_DETAIL, data["detail"]) self.assertIsValidEvent(data)