def test_rollback_all(self, mock_helper):
"""CLI - Lambda rollback all"""
mock_helper.return_value = True
funcs = [
'alert', 'alert_merger', 'apps', 'athena', 'classifier', 'rule',
'rule_promo', 'scheduled_queries', 'threat_intel_downloader'
]
assert_equal(
rollback.RollbackCommand.handler(
MockOptions(None, funcs),
MockCLIConfig(config=basic_streamalert_config())), True)
mock_helper.assert_has_calls([
mock.call(mock.ANY, 'unit-test_streamalert_alert_processor'),
mock.call(mock.ANY, 'unit-test_streamalert_alert_merger'),
mock.call(mock.ANY,
'unit-test_corp_box_admin_events_box_collector_app'),
mock.call(mock.ANY,
'unit-test_corp_duo_admin_duo_admin_collector_app'),
mock.call(mock.ANY, 'unit-test_streamalert_athena_partitioner'),
mock.call(mock.ANY, 'unit-test_corp_streamalert_classifier'),
mock.call(mock.ANY, 'unit-test_prod_streamalert_classifier'),
mock.call(mock.ANY, 'unit-test_streamalert_rules_engine'),
mock.call(mock.ANY, 'unit-test_streamalert_rule_promotion'),
mock.call(mock.ANY,
'unit-test_streamalert_scheduled_queries_runner'),
mock.call(mock.ANY,
'unit-test_streamalert_threat_intel_downloader')
])
def test_rollback_subset(self, mock_helper):
"""CLI - Lambda rollback apps and rule"""
mock_helper.return_value = True
assert_equal(
rollback.rollback_handler(
MockOptions(None, ['apps', 'rule']),
MockCLIConfig(config=basic_streamalert_config())), True)
mock_helper.assert_has_calls([
mock.call(mock.ANY,
'unit-testing_corp_box_admin_events_box_collector_app'),
mock.call(mock.ANY,
'unit-testing_corp_duo_admin_duo_admin_collector_app'),
mock.call(mock.ANY, 'unit-testing_streamalert_rules_engine')
])
def test_rollback_all(self, mock_helper):
"""CLI - Lambda rollback all"""
mock_helper.return_value = True
assert_equal(
rollback.rollback_handler(
MockOptions(None, ['all']),
MockCLIConfig(config=basic_streamalert_config())), True)
mock_helper.assert_has_calls([
mock.call(mock.ANY, 'unit-testing_streamalert_alert_processor'),
mock.call(mock.ANY, 'unit-testing_streamalert_alert_merger'),
mock.call(mock.ANY,
'unit-testing_corp_box_admin_events_box_collector_app'),
mock.call(mock.ANY,
'unit-testing_corp_duo_admin_duo_admin_collector_app'),
mock.call(mock.ANY,
'unit-testing_streamalert_athena_partition_refresh'),
mock.call(mock.ANY, 'unit-testing_corp_streamalert_classifier'),
mock.call(mock.ANY, 'unit-testing_prod_streamalert_classifier'),
mock.call(mock.ANY, 'unit-testing_streamalert_rules_engine'),
mock.call(mock.ANY,
'unit-testing_streamalert_threat_intel_downloader')
])
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))