def test_repair_hive_table_unknown_bucket(self, mock_logging):
"""Athena - Repair Hive Table - Unknown Bucket"""
self.client.athena_client = MockAthenaClient(result_state='SUCCEEDED')
# This bucket is not in our `repair_hive_table` config map
self.client.repair_hive_table({'my-test.result.bucket'})
assert_true(mock_logging.warning.called)
def test_repair_hive_table_failed_refresh(self, mock_logging):
"""Athena - Repair Hive Table - Failed Refresh"""
self.client.athena_client = MockAthenaClient(result_state='FAILED')
# This bucket is not in our `repair_hive_table` config map
self.client.repair_hive_table({'unit-testing.streamalerts'})
assert_true(mock_logging.error.called)
def test_add_hive_partition(self, mock_logging):
"""Athena - Add Hive Partition"""
query_result = [{'Repair: added data to metastore:foobar'},
{'Repair: added data to metastore:foobaz'}]
self.client.athena_client = MockAthenaClient(results=query_result)
result = self.client.add_hive_partition({
'unit-testing.streamalerts':
set([
'alerts/dt=2017-08-26-14/rule_name_alerts-1304134918401.json',
'alerts/dt=2017-08-27-14/rule_name_alerts-1304134918401.json'
]),
'unit-testing.streamalert.data':
set([
'log_type_1/2017/08/26/14/test-data-11111-22222-33333.snappy',
'log_type_2/2017/08/26/14/test-data-11111-22222-33333.snappy',
'log_type_2/2017/08/26/15/test-data-11111-22222-33333.snappy',
'log_type_2/2017/08/26/16/test-data-11111-22222-33333.snappy',
'log_type_3/2017/08/26/14/test-data-11111-22222-33333.snappy',
'log_type_1/2017/08/26/11/test-data-11111-22222-33333.snappy'
]),
'test-bucket-with-data':
set([
'2017/08/26/14/rule_name_alerts-1304134918401.json',
'2017/08/28/14/rule_name_alerts-1304134918401.json',
'2017/07/30/14/rule_name_alerts-1304134918401.json'
])
})
assert_true(mock_logging.info.called)
assert_true(result)
def test_repair_hive_table(self, mock_logging):
"""Athena - Repair Hive Table"""
query_result = [{'Status': 'SUCCEEDED'}]
self.client.athena_client = MockAthenaClient(results=query_result)
self.client.repair_hive_table({'unit-testing.streamalerts'})
assert_true(mock_logging.info.called)
def test_check_table_exists_invalid(self, mock_logging):
"""Athena - Check Table Exists - Does Not Exist"""
query_result = None
self.client.athena_client = MockAthenaClient(results=query_result)
result = self.client.check_table_exists('unit-test')
assert_false(result)
assert_true(mock_logging.info.called)
def test_run_athena_query_async(self):
"""Athena - Run Athena Query - Async Call"""
query_result = []
self.client.athena_client = MockAthenaClient(results=query_result)
query_success, _ = self.client.run_athena_query(
query='SHOW DATABASES;', async=True)
assert_true(query_success)
def setup(self, boto_patch):
"""Setup the AthenaRefresher tests"""
self.mock_sqs = mock_sqs()
self.mock_sqs.start()
boto_patch.client.return_value = MockAthenaClient()
sqs = boto3.resource('sqs')
name = StreamAlertSQSClient.DEFAULT_QUEUE_NAME.format('unit-testing')
self.queue = sqs.create_queue(QueueName=name)
self._refresher = AthenaRefresher()
def test_run_athena_query(self):
"""Athena - Run Athena Query - Normal Result"""
self.client.athena_client = MockAthenaClient()
query_success, query_results = self.client.run_athena_query(
query='SHOW DATABASES;'
)
assert_true(query_success)
assert_equal(query_results['ResultSet']['Rows'], [{'Data': [{'test': 'test'}]}])
def test_check_table_exists(self):
"""Athena - Check Table Exists"""
query_result = [{'alerts': True}]
self.client.athena_client = MockAthenaClient(results=query_result)
result = self.client.check_table_exists('unit-test')
assert_true(result)
generated_results_key = 'unit-testing/{}'.format(
datetime.now().strftime('%Y/%m/%d'))
assert_equal(self.client.athena_results_key, generated_results_key)
def test_run_athena_query_error(self, mock_logging):
"""Athena - Run Athena Query - Error Result"""
self.client.athena_client = MockAthenaClient(results=None,
result_state='FAILED')
query_success, query_results = self.client.run_athena_query(
query='SHOW DATABASES;')
assert_true(mock_logging.error.called)
assert_false(query_success)
assert_equal(query_results, {})
def test_run_athena_query_empty(self, mock_logging):
"""Athena - Run Athena Query - Empty Result"""
query_result = None
self.client.athena_client = MockAthenaClient(results=query_result)
query_success, query_results = self.client.run_athena_query(
query='SHOW DATABASES;')
assert_true(query_success)
assert_equal(query_results['ResultSet']['Rows'], [])
assert_true(mock_logging.debug.called)
def test_add_hive_partition_unknown_bucket(self, mock_logging):
"""Athena - Add Hive Partition - Unknown Bucket"""
self.client.athena_client = MockAthenaClient(results=[])
result = self.client.add_hive_partition({
'bucket-not-in-config.streamalerts': set([
'alerts/dt=2017-08-26-14/rule_name_alerts-1304134918401.json',
'alerts/dt=2017-08-27-14/rule_name_alerts-1304134918401.json',
])
})
assert_true(mock_logging.error.called)
assert_false(result)
def test_add_hive_partition_unexpected_s3_key(self, mock_logging):
"""Athena - Add Hive Partition - Unexpected S3 Key"""
self.client.athena_client = MockAthenaClient(results=[])
result = self.client.add_hive_partition({
'unit-testing.streamalerts':
set(['a/pattern/that/does/not-match']),
'test-bucket-with-data':
set(['another/pattern/that/does/not-match'])
})
assert_true(mock_logging.error.called)
assert_false(result)
def test_rebuild_partitions():
"""CLI - Athena rebuild partitions helper"""
with patch('streamalert.shared.athena.boto3') as mock_athena:
mock_show_partitions_result = [{
'Data': [{
'VarCharValue': 'dt=2019-12-04-05'
}]
}, {
'Data': [{
'VarCharValue': 'dt=2019-12-03-22'
}]
}, {
'Data': [{
'VarCharValue': 'dt=2019-12-03-23'
}]
}, {
'Data': [{
'VarCharValue': 'dt=2019-12-03-20'
}]
}, {
'Data': [{
'VarCharValue': 'dt=2019-12-04-01'
}]
}]
mock_show_table_result = []
mock_athena.client.side_effect = [
MockAthenaClient(results=mock_show_partitions_result),
MockAthenaClient(results=mock_show_table_result)
]
config = MockCLIConfig(config=athena_cli_basic_config())
table = 'unit_my_test'
bucket = 'bucket'
assert_true(handler.rebuild_partitions(table, bucket, config))
def _mock_boto(name):
"""Hack to allow mocking boto3.client with moto and our own class"""
if name == 'athena':
return MockAthenaClient()
return client(name)
def setup(self, boto_patch):
"""Setup the AthenaRefresher tests"""
boto_patch.client.return_value = MockAthenaClient()
self._refresher = AthenaRefresher()
def test_check_database_exists(self):
"""Athena - Check Database Exists"""
query_result = [{'streamalert': True}]
self.client.athena_client = MockAthenaClient(results=query_result)
assert_true(self.client.check_database_exists())
def test_check_database_exists_invalid(self):
"""Athena - Check Database Exists - Does Not Exist"""
query_result = None
self.client.athena_client = MockAthenaClient(results=query_result)
assert_false(self.client.check_database_exists())
def setup(self, boto_patch):
"""Setup the AthenaPartitioner tests"""
boto_patch.client.return_value = MockAthenaClient()
self._partitioner = AthenaPartitioner()
class TestStreamAlertAthenaClient(object):
"""Test class for StreamAlertAthenaClient"""
@patch.dict(os.environ, {'AWS_DEFAULT_REGION': 'us-west-1'})
@patch('boto3.client', Mock(side_effect=lambda c: MockAthenaClient()))
def setup(self):
"""Setup the StreamAlertAthenaClient tests"""
self._db_name = 'test_database'
config = load_config('tests/unit/conf/')
prefix = config['global']['account']['prefix']
self.client = StreamAlertAthenaClient(
self._db_name, 's3://{}.streamalert.athena-results'.format(prefix),
'unit-testing')
@patch('stream_alert.shared.athena.datetime')
def test_init_fix_bucket_path(self, date_mock):
"""Athena - Fix Bucket Path"""
date_now = datetime.utcnow()
date_mock.utcnow.return_value = date_now
date_format = date_now.strftime('%Y/%m/%d')
expected_path = 's3://test.streamalert.athena-results/unit-testing/{}'.format(
date_format)
with patch.dict(os.environ, {'AWS_DEFAULT_REGION': 'us-west-1'}):
client = StreamAlertAthenaClient(
self._db_name, 'test.streamalert.athena-results',
'unit-testing')
assert_equal(client._s3_results_path, expected_path)
def test_unique_values_from_query(self):
"""Athena - Unique Values from Query"""
query = {
'ResultSet': {
'Rows': [
{
'Data': [{
'VarCharValue': 'foobar'
}]
},
{
'Data': [{
'VarCharValue': 'barfoo'
}]
},
{
'Data': [{
'VarCharValue': 'barfoo'
}]
},
{
'Data': [{
'VarCharValue': 'foobarbaz'
}]
},
]
}
}
expected_result = {'foobar', 'barfoo', 'foobarbaz'}
result = self.client._unique_values_from_query(query)
assert_items_equal(result, expected_result)
def test_check_database_exists(self):
"""Athena - Check Database Exists"""
self.client._client.results = [{
'Data': [{
'VarCharValue': self._db_name
}]
}]
assert_true(self.client.check_database_exists())
def test_check_database_exists_invalid(self):
"""Athena - Check Database Exists - Does Not Exist"""
self.client._client.results = None
assert_false(self.client.check_database_exists())
def test_check_table_exists(self):
"""Athena - Check Table Exists"""
self.client._client.results = [{
'Data': [{
'VarCharValue': 'test_table'
}]
}]
assert_true(self.client.check_table_exists('test_table'))
def test_check_table_exists_invalid(self):
"""Athena - Check Table Exists - Does Not Exist"""
self.client._client.results = None
assert_false(self.client.check_table_exists('test_table'))
def test_get_table_partitions(self):
"""Athena - Get Table Partitions"""
self.client._client.results = [
{
'Data': [{
'VarCharValue': 'dt=2018-12-10-10'
}]
},
{
'Data': [{
'VarCharValue': 'dt=2018-12-09-10'
}]
},
{
'Data': [{
'VarCharValue': 'dt=2018-12-09-10'
}]
},
{
'Data': [{
'VarCharValue': 'dt=2018-12-11-10'
}]
},
]
expected_result = {
'dt=2018-12-10-10', 'dt=2018-12-09-10', 'dt=2018-12-11-10'
}
result = self.client.get_table_partitions('test_table')
assert_items_equal(result, expected_result)
def test_get_table_partitions_error(self):
"""Athena - Get Table Partitions, Exception"""
self.client._client.raise_exception = True
result = self.client.get_table_partitions('test_table')
assert_equal(result, None)
def test_drop_table(self):
"""Athena - Drop Table, Success"""
assert_true(self.client.drop_table('test_table'))
def test_drop_table_failure(self):
"""Athena - Drop Table, Failure"""
self.client._client.raise_exception = True
assert_false(self.client.drop_table('test_table'))
@patch('stream_alert.shared.athena.StreamAlertAthenaClient.drop_table')
def test_drop_all_tables(self, drop_table_mock):
"""Athena - Drop All Tables, Success"""
self.client._client.results = [
{
'Data': [{
'VarCharValue': 'table_01'
}]
},
{
'Data': [{
'VarCharValue': 'table_02'
}]
},
{
'Data': [{
'VarCharValue': 'table_02'
}]
},
]
assert_true(self.client.drop_all_tables())
assert_equal(drop_table_mock.call_count, 2)
@patch('stream_alert.shared.athena.StreamAlertAthenaClient.drop_table')
def test_drop_all_tables_failure(self, drop_table_mock):
"""Athena - Drop All Tables, Failure"""
self.client._client.results = [
{
'Data': [{
'VarCharValue': 'table_01'
}]
},
{
'Data': [{
'VarCharValue': 'table_02'
}]
},
{
'Data': [{
'VarCharValue': 'table_03'
}]
},
]
drop_table_mock.side_effect = [True, True, False]
assert_false(self.client.drop_all_tables())
def test_drop_all_tables_exception(self):
"""Athena - Drop All Tables, Exception"""
self.client._client.raise_exception = True
assert_false(self.client.drop_all_tables())
@patch('logging.Logger.exception')
def test_execute_query(self, log_mock):
"""Athena - Execute Query"""
self.client._client.raise_exception = True
self.client._execute_query('BAD SQL')
log_mock.assert_called_with('Athena query failed')
def test_execute_and_wait(self):
"""Athena - Execute and Wait"""
self.client._client.results = [
{
'Data': [{
'VarCharValue': 'result'
}]
},
]
result = self.client._execute_and_wait('SQL query')
assert_true(result in self.client._client.query_executions)
@patch('logging.Logger.error')
def test_execute_and_wait_failed(self, log_mock):
"""Athena - Execute and Wait, Failed"""
query = 'SQL query'
self.client._client.result_state = 'FAILED'
self.client._execute_and_wait(query)
log_mock.assert_called_with('Athena query failed:\n%s', query)
def test_query_result_paginator(self):
"""Athena - Query Result Paginator"""
data = {'Data': [{'VarCharValue': 'result'}]}
self.client._client.results = [
data,
]
items = list(self.client.query_result_paginator('test query'))
assert_items_equal(items, [{'ResultSet': {'Rows': [data]}}] * 4)
def test_query_result_paginator_error(self):
"""Athena - Query Result Paginator, Exception"""
self.client._client.raise_exception = True
assert_equal(list(self.client.query_result_paginator('test query')),
list())
def test_run_async_query(self):
"""Athena - Run Async Query, Success"""
assert_true(self.client.run_async_query('test query'))
def test_run_async_query_failure(self):
"""Athena - Run Async Query, Failure"""
self.client._client.raise_exception = True
assert_false(self.client.run_async_query('test query'))
