def _add_fake_stats(self):
"""Helper function to add a few fake alert statistic objects"""
time = datetime(year=2000, month=1, day=1, hour=1, second=1)
self.promoter._staging_stats.update({
'test_rule': StagingStatistic(time, time + timedelta(days=2), time, 'test_rule'),
'test_rule_2': StagingStatistic(time, time + timedelta(days=2), time, 'test_rule_2')
})
def test_comp(self):
"""StagingStatistic - Comparison"""
self.statistic.alert_count = 200
second_stat = StagingStatistic(staged_at='fake_staged_at_time',
staged_until='fake_staged_until_time',
current_time='fake_current_time',
rule='test_rule')
second_stat.alert_count = 100
assert_equal(self.statistic > second_stat, True)
def _get_fake_stats(count=2):
"""Helper function to return fake StagingStatistics"""
stage_time = datetime(year=2000, month=1, day=1, hour=1, minute=1, second=1)
for i in range(count):
stat = StagingStatistic(
staged_at=stage_time,
staged_until=stage_time + timedelta(days=2),
current_time=stage_time + timedelta(days=1),
rule='test_rule_{}'.format(i)
)
stat.alert_count = i + 1
yield stat
def _update_alert_count(self):
"""Transform Athena query results into alert counts for rules_engine
Args:
query (str): Athena query to run and wait for results
Returns:
dict: Representation of alert counts, where key is the rule name
and value is the alert count (int) since this rule was staged
"""
query = StagingStatistic.construct_compound_count_query(
self._staging_stats.values())
LOGGER.debug('Running compound query for alert count: \'%s\'', query)
for page, results in enumerate(
self._athena_client.query_result_paginator(query)):
for i, row in enumerate(results['ResultSet']['Rows']):
if page == 0 and i == 0: # skip header row included in first page only
continue
row_values = [data.values()[0] for data in row['Data']]
rule_name, alert_count = row_values[0], int(row_values[1])
LOGGER.debug('Found %d alerts for rule \'%s\'', alert_count,
rule_name)
self._staging_stats[rule_name].alert_count = alert_count
def _get_staging_info(self):
"""Query the Rule table for rule staging info needed to count each rule's alerts
Example of rule metadata returned by RuleTable.remote_rule_info():
{
'example_rule_name':
{
'Staged': True
'StagedAt': datetime.datetime object,
'StagedUntil': '2018-04-21T02:23:13.332223Z'
}
}
"""
for rule in sorted(self._rule_table.remote_rule_info):
info = self._rule_table.remote_rule_info[rule]
# If the rule is not staged, do not get stats on it
if not info['Staged']:
continue
self._staging_stats[rule] = StagingStatistic(
info['StagedAt'],
info['StagedUntil'],
self._current_time,
rule
)
return len(self._staging_stats) != 0
def test_construct_compound_count_query(self):
"""StagingStatistic - Construct Compound Count Query"""
query = StagingStatistic.construct_compound_count_query([self.statistic, self.statistic])
expected_query = ("SELECT rule_name, count(*) AS count "
"FROM alerts WHERE "
"(dt >= '2000-01-01-01' AND rule_name = 'test_rule') OR "
"(dt >= '2000-01-01-01' AND rule_name = 'test_rule') "
"GROUP BY rule_name")
assert_equal(query, expected_query)
def setup(self):
"""StagingStatistic - Setup"""
# pylint: disable=attribute-defined-outside-init
stage_time = datetime(year=2000, month=1, day=1, hour=1, minute=1, second=1)
self.statistic = StagingStatistic(
staged_at=stage_time,
staged_until=stage_time + timedelta(days=2),
current_time=stage_time + timedelta(days=1),
rule='test_rule'
)
def test_construct_compound_count_query(self):
"""StagingStatistic - Construct Compound Count Query"""
query = StagingStatistic.construct_compound_count_query(
[self.statistic, self.statistic])
expected_query = (
"SELECT 'test_rule' AS rule_name, count(*) AS count "
"FROM alerts WHERE dt >= '2000-01-01-01' AND "
"rule_name = 'test_rule' UNION ALL SELECT 'test_rule' "
"AS rule_name, count(*) AS count FROM alerts WHERE "
"dt >= '2000-01-01-01' AND rule_name = 'test_rule'")
assert_equal(query, expected_query)