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)