def test_load_from_config(self):
"""FirehoseClient - Load From Config"""
with patch('boto3.client'): # patch to speed up unit tests slightly
client = FirehoseClient.load_from_config(
prefix='unit-test',
firehose_config={'enabled': True},
log_sources=None)
assert_equal(isinstance(client, FirehoseClient), True)
def __init__(self):
# Create some objects to be cached if they have not already been created
Classifier._config = Classifier._config or config.load_config(validate=True)
Classifier._firehose_client = (
Classifier._firehose_client or FirehoseClient.load_from_config(
prefix=self.config['global']['account']['prefix'],
firehose_config=self.config['global'].get('infrastructure', {}).get('firehose', {}),
log_sources=self.config['logs']
)
)
Classifier._sqs_client = Classifier._sqs_client or SQSClient()
# Setup the normalization logic
Normalizer.load_from_config(self.config)
self._cluster = os.environ['CLUSTER']
self._payloads = []
self._failed_record_count = 0
self._processed_size = 0
def test_send_no_prefixing(self, send_batch_mock):
"""FirehoseClient - Send, No Prefixing"""
FirehoseClient._ENABLED_LOGS = {
'log_type_01_sub_type_01': 'log_type_01:sub_type_01'
}
expected_batch = [
'{"unit_key_01":1,"unit_key_02":"test"}\n',
'{"unit_key_01":2,"unit_key_02":"test"}\n'
]
client = FirehoseClient.load_from_config(prefix='unit-test',
firehose_config={
'enabled': True,
'use_prefix': False
},
log_sources=None)
client.send(self._sample_payloads)
send_batch_mock.assert_called_with(
'streamalert_log_type_01_sub_type_01', expected_batch,
'classifier')
def test_send_long_log_name(self, send_batch_mock):
"""FirehoseClient - Send data when the log name is very long"""
FirehoseClient._ENABLED_LOGS = {
'very_very_very_long_log_stream_name_abcdefg_hijklmn_70_characters_long':
{}
}
expected_batch = [
'{"unit_key_01":1,"unit_key_02":"test"}\n',
'{"unit_key_01":2,"unit_key_02":"test"}\n'
]
client = FirehoseClient.load_from_config(prefix='unit-test',
firehose_config={
'enabled': True,
'use_prefix': False
},
log_sources=None)
client.send(self._sample_payloads_long_log_name)
send_batch_mock.assert_called_with(
'streamalert_very_very_very_long_log_stream_name_abcdefg_7c88167b',
expected_batch, 'classifier')
def test_load_from_config_disabled(self):
"""FirehoseClient - Load From Config, Disabled"""
client = FirehoseClient.load_from_config(prefix='unit-test',
firehose_config={},
log_sources=None)
assert_equal(client, None)