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( list(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 = [list(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 test_comp_no_alert_count(self): """StagingStatistic - Comparison when alert_count is default value""" # 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, False) self.statistic._current_time += timedelta(days=2, hours=10) expected_string = '''\u25E6 test_rule - Staged At: 2000-01-01 01:01:01 UTC - Staged Until: 2000-01-03 01:01:01 UTC - Time Past Staging: 1d 10h 0m - Alert Count: unknown - Alert Info: n/a''' assert_equal(str(self.statistic), expected_string)
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 _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