def test_load_from_raw_record_app(self):
"""StreamPayload - Load from Raw Record, StreamAlertApp"""
record = {'stream_alert_app': 'test_app'}
with patch.object(RegisterInput, 'load_for_service') as load_mock:
StreamPayload.load_from_raw_record(record)
load_mock.assert_called_with('stream_alert_app', 'test_app',
record)
def test_load_from_raw_record_kinesis(self):
"""StreamPayload - Load from Raw Record, Kinesis"""
record = {
'kinesis': {},
'eventSourceARN':
'arn:aws:kinesis:us-east-1:123456789012:stream/test_stream_name'
}
with patch.object(RegisterInput, 'load_for_service') as load_mock:
StreamPayload.load_from_raw_record(record)
load_mock.assert_called_with('kinesis', 'test_stream_name', record)
def test_load_from_raw_record_sns(self):
"""StreamPayload - Load from Raw Record, SNS"""
record = {
'Sns': {
'TopicArn':
'arn:aws:sns:us-east-1:123456789012:test_topic_arn',
}
}
with patch.object(RegisterInput, 'load_for_service') as load_mock:
StreamPayload.load_from_raw_record(record)
load_mock.assert_called_with('sns', 'test_topic_arn', record)
def test_load_from_raw_record_sns_s3(self):
"""StreamPayload - Load from Raw Record, SNS S3 Event"""
s3_record = {'s3': {'bucket': {'name': 'test_bucket_name'}}}
record = {
'Sns': {
'Type': 'Notification',
'Subject': 'Amazon S3 Notification',
'TopicArn':
'arn:aws:sns:us-east-1:123456789012:test_topic_arn',
'Message': json.dumps({'Records': [s3_record]})
}
}
with patch.object(RegisterInput, 'load_for_service') as load_mock:
StreamPayload.load_from_raw_record(record)
load_mock.assert_called_with('s3', 'test_bucket_name', s3_record)
def run(self, records):
"""Run classificaiton of the records in the Lambda input
Args:
records (list): An list of records received by Lambda
"""
LOGGER.debug('Number of incoming records: %d', len(records))
if not records:
return
for input_record in records:
# Get the service and entity from the payload
payload = StreamPayload.load_from_raw_record(input_record)
if not payload:
self._log_bad_records(input_record, 1)
continue
self._classify_payload(payload)
self._log_metrics()
# Send records to SQS before sending to Firehose
self.sqs.send(self._payloads)
# Send the data to firehose for historical retention
if self.data_retention_enabled:
self.firehose.send(self._payloads)
return self._payloads
def setup(self):
"""StreamPayload - Setup"""
# pylint: disable=abstract-class-instantiated,attribute-defined-outside-init
self._resource = 'foobar'
self._record = {'key': 'value'}
self._payload = StreamPayload(self._resource, self._record)
def test_load_from_raw_record_s3(self):
"""StreamPayload - Load from Raw Record, S3"""
record = {'s3': {'bucket': {'name': 'test_bucket_name'}}}
with patch.object(RegisterInput, 'load_for_service') as load_mock:
StreamPayload.load_from_raw_record(record)
load_mock.assert_called_with('s3', 'test_bucket_name', record)