<commit_before># -*- coding: utf-8 -*-
import contextlib
import copy
import datetime
import json
import threading

import elasticsearch
import mock
import pytest
from elasticsearch.exceptions import ElasticsearchException

from elastalert.enhancements import BaseEnhancement
from elastalert.enhancements import DropMatchException
from elastalert.kibana import dashboard_temp
from elastalert.util import dt_to_ts
from elastalert.util import dt_to_unix
from elastalert.util import dt_to_unixms
from elastalert.util import EAException
from elastalert.util import ts_to_dt
from elastalert.util import unix_to_dt


START_TIMESTAMP = '2014-09-26T12:34:45Z'
END_TIMESTAMP = '2014-09-27T12:34:45Z'
START = ts_to_dt(START_TIMESTAMP)
END = ts_to_dt(END_TIMESTAMP)


def _set_hits(ea_inst, hits):
    res = {'hits': {'total': len(hits), 'hits': hits}}
    ea_inst.client_es.return_value = res


def generate_hits(timestamps, **kwargs):
    hits = []
    id_iter = xrange(len(timestamps)).__iter__()
    for ts in timestamps:
        data = {'_id': 'id' + str(id_iter.next()),
                '_source': {'@timestamp': ts},
                '_type': 'logs',
                '_index': 'idx'}
        for key, item in kwargs.iteritems():
            data['_source'][key] = item
        # emulate process_hits(), add metadata to _source
        for field in ['_id', '_type', '_index']:
            data['_source'][field] = data[field]
        hits.append(data)
    return {'hits': {'total': len(hits), 'hits': hits}}


def assert_alerts(ea_inst, calls):
    """ Takes a list of lists of timestamps. Asserts that an alert was called for each list, containing those timestamps. """
    assert ea_inst.rules[0]['alert'][0].alert.call_count == len(calls)
    for call_num, call_args in enumerate(ea_inst.rules[0]['alert'][0].alert.call_args_list):
        assert not any([match['@timestamp'] not in calls[call_num] for match in call_args[0][0]])
        assert len(call_args[0][0]) == len(calls[call_num])


def test_starttime(ea):
    invalid = ['2014-13-13',
               '2014-11-24T30:00:00',
               'Not A Timestamp']
    for ts in invalid:
        with pytest.raises((TypeError, ValueError)):
            ts_to_dt(ts)


def test_init_rule(ea):
    # Simulate state of a rule just loaded from a file
    ea.rules[0]['minimum_starttime'] = datetime.datetime.now()
    new_rule = copy.copy(ea.rules[0])
    map(new_rule.pop, ['agg_matches', 'current_aggregate_id', 'processed_hits', 'minimum_starttime'])

    # Properties are copied from ea.rules[0]
    ea.rules[0]['starttime'] = '2014-01-02T00:11:22'
    ea.rules[0]['processed_hits'] = ['abcdefg']
    new_rule = ea.init_rule(new_rule, False)
    for prop in ['starttime', 'agg_matches', 'current_aggregate_id', 'processed_hits', 'minimum_starttime']:
        assert new_rule[prop] == ea.rules[0][prop]

    # Properties are fresh
    new_rule = ea.init_rule(new_rule, True)
    new_rule.pop('starttime')
    assert 'starttime' not in new_rule
    assert new_rule['processed_hits'] == {}


def test_query(ea):
    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.run_query(ea.rules[0], START, END)
    ea.current_es.search.assert_called_with(body={'query': {'filtered': {'filter': {'bool': {'must': [{'range': {'@timestamp': {'lte': END_TIMESTAMP, 'gt': START_TIMESTAMP}}}]}}}}, 'sort': [{'@timestamp': {'order': 'asc'}}]}, index='idx', _source_include=['@timestamp'], ignore_unavailable=True, size=ea.rules[0]['max_query_size'], scroll=ea.conf['scroll_keepalive'])


def test_query_with_fields(ea):
    ea.rules[0]['_source_enabled'] = False
    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.run_query(ea.rules[0], START, END)
    ea.current_es.search.assert_called_with(body={'query': {'filtered': {'filter': {'bool': {'must': [{'range': {'@timestamp': {'lte': END_TIMESTAMP, 'gt': START_TIMESTAMP}}}]}}}}, 'sort': [{'@timestamp': {'order': 'asc'}}], 'fields': ['@timestamp']}, index='idx', ignore_unavailable=True, size=ea.rules[0]['max_query_size'], scroll=ea.conf['scroll_keepalive'])


def test_query_with_unix(ea):
    ea.rules[0]['timestamp_type'] = 'unix'
    ea.rules[0]['dt_to_ts'] = dt_to_unix
    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.run_query(ea.rules[0], START, END)
    start_unix = dt_to_unix(START)
    end_unix = dt_to_unix(END)
    ea.current_es.search.assert_called_with(body={'query': {'filtered': {'filter': {'bool': {'must': [{'range': {'@timestamp': {'lte': end_unix, 'gt': start_unix}}}]}}}}, 'sort': [{'@timestamp': {'order': 'asc'}}]}, index='idx', _source_include=['@timestamp'], ignore_unavailable=True, size=ea.rules[0]['max_query_size'], scroll=ea.conf['scroll_keepalive'])


def test_query_with_unixms(ea):
    ea.rules[0]['timestamp_type'] = 'unixms'
    ea.rules[0]['dt_to_ts'] = dt_to_unixms
    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.run_query(ea.rules[0], START, END)
    start_unix = dt_to_unixms(START)
    end_unix = dt_to_unixms(END)
    ea.current_es.search.assert_called_with(body={'query': {'filtered': {'filter': {'bool': {'must': [{'range': {'@timestamp': {'lte': end_unix, 'gt': start_unix}}}]}}}}, 'sort': [{'@timestamp': {'order': 'asc'}}]}, index='idx', _source_include=['@timestamp'], ignore_unavailable=True, size=ea.rules[0]['max_query_size'], scroll=ea.conf['scroll_keepalive'])


def test_no_hits(ea):
    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.run_query(ea.rules[0], START, END)
    assert ea.rules[0]['type'].add_data.call_count == 0


def test_no_terms_hits(ea):
    ea.rules[0]['use_terms_query'] = True
    ea.rules[0]['query_key'] = 'QWERTY'
    ea.rules[0]['doc_type'] = 'uiop'
    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.run_query(ea.rules[0], START, END)
    assert ea.rules[0]['type'].add_terms_data.call_count == 0


def test_some_hits(ea):
    hits = generate_hits([START_TIMESTAMP, END_TIMESTAMP])
    hits_dt = generate_hits([START, END])
    ea.current_es.search.return_value = hits
    ea.run_query(ea.rules[0], START, END)
    assert ea.rules[0]['type'].add_data.call_count == 1
    ea.rules[0]['type'].add_data.assert_called_with([x['_source'] for x in hits_dt['hits']['hits']])


def test_some_hits_unix(ea):
    ea.rules[0]['timestamp_type'] = 'unix'
    ea.rules[0]['dt_to_ts'] = dt_to_unix
    ea.rules[0]['ts_to_dt'] = unix_to_dt
    hits = generate_hits([dt_to_unix(START), dt_to_unix(END)])
    hits_dt = generate_hits([START, END])
    ea.current_es.search.return_value = copy.deepcopy(hits)
    ea.run_query(ea.rules[0], START, END)
    assert ea.rules[0]['type'].add_data.call_count == 1
    ea.rules[0]['type'].add_data.assert_called_with([x['_source'] for x in hits_dt['hits']['hits']])


def _duplicate_hits_generator(timestamps, **kwargs):
    """Generator repeatedly returns identical hits dictionaries
    """
    while True:
        yield generate_hits(timestamps, **kwargs)


def test_duplicate_timestamps(ea):
    ea.current_es.search.side_effect = _duplicate_hits_generator([START_TIMESTAMP] * 3, blah='duplicate')
    ea.run_query(ea.rules[0], START, ts_to_dt('2014-01-01T00:00:00Z'))

    assert len(ea.rules[0]['type'].add_data.call_args_list[0][0][0]) == 3
    assert ea.rules[0]['type'].add_data.call_count == 1

    # Run the query again, duplicates will be removed and not added
    ea.run_query(ea.rules[0], ts_to_dt('2014-01-01T00:00:00Z'), END)
    assert ea.rules[0]['type'].add_data.call_count == 1


def test_match(ea):
    hits = generate_hits([START_TIMESTAMP, END_TIMESTAMP])
    ea.current_es.search.return_value = hits
    ea.rules[0]['type'].matches = [{'@timestamp': END}]
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        ea.run_rule(ea.rules[0], END, START)

    ea.rules[0]['alert'][0].alert.called_with({'@timestamp': END_TIMESTAMP})
    assert ea.rules[0]['alert'][0].alert.call_count == 1


def test_run_rule_calls_garbage_collect(ea):
    start_time = '2014-09-26T00:00:00Z'
    end_time = '2014-09-26T12:00:00Z'
    ea.buffer_time = datetime.timedelta(hours=1)
    ea.run_every = datetime.timedelta(hours=1)
    with contextlib.nested(mock.patch.object(ea.rules[0]['type'], 'garbage_collect'),
                           mock.patch.object(ea, 'run_query')) as (mock_gc, mock_get_hits):
        ea.run_rule(ea.rules[0], ts_to_dt(end_time), ts_to_dt(start_time))

    # Running elastalert every hour for 12 hours, we should see self.garbage_collect called 12 times.
    assert mock_gc.call_count == 12

    # The calls should be spaced 1 hour apart
    expected_calls = [ts_to_dt(start_time) + datetime.timedelta(hours=i) for i in range(1, 13)]
    for e in expected_calls:
        mock_gc.assert_any_call(e)


def run_rule_query_exception(ea, mock_es):
    with mock.patch('elastalert.elastalert.elasticsearch_client') as mock_es_init:
        mock_es_init.return_value = mock_es
        ea.run_rule(ea.rules[0], END, START)

    # Assert neither add_data nor garbage_collect were called
    # and that starttime did not change
    assert ea.rules[0].get('starttime') == START
    assert ea.rules[0]['type'].add_data.call_count == 0
    assert ea.rules[0]['type'].garbage_collect.call_count == 0
    assert ea.rules[0]['type'].add_count_data.call_count == 0


def test_query_exception(ea):
    mock_es = mock.Mock()
    mock_es.search.side_effect = ElasticsearchException
    run_rule_query_exception(ea, mock_es)


def test_query_exception_count_query(ea):
    ea.rules[0]['use_count_query'] = True
    ea.rules[0]['doc_type'] = 'blahblahblahblah'
    mock_es = mock.Mock()
    mock_es.count.side_effect = ElasticsearchException
    run_rule_query_exception(ea, mock_es)


def test_match_with_module(ea):
    mod = BaseEnhancement(ea.rules[0])
    mod.process = mock.Mock()
    ea.rules[0]['match_enhancements'] = [mod]
    test_match(ea)
    mod.process.assert_called_with({'@timestamp': END})


def test_match_with_module_with_agg(ea):
    mod = BaseEnhancement(ea.rules[0])
    mod.process = mock.Mock()
    ea.rules[0]['match_enhancements'] = [mod]
    ea.rules[0]['aggregation'] = datetime.timedelta(minutes=15)
    hits = generate_hits([START_TIMESTAMP, END_TIMESTAMP])
    ea.current_es.search.return_value = hits
    ea.rules[0]['type'].matches = [{'@timestamp': END}]
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        ea.run_rule(ea.rules[0], END, START)
    assert mod.process.call_count == 0


def test_match_with_enhancements_first(ea):
    mod = BaseEnhancement(ea.rules[0])
    mod.process = mock.Mock()
    ea.rules[0]['match_enhancements'] = [mod]
    ea.rules[0]['aggregation'] = datetime.timedelta(minutes=15)
    ea.rules[0]['run_enhancements_first'] = True
    hits = generate_hits([START_TIMESTAMP, END_TIMESTAMP])
    ea.current_es.search.return_value = hits
    ea.rules[0]['type'].matches = [{'@timestamp': END}]
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        with mock.patch.object(ea, 'add_aggregated_alert') as add_alert:
            ea.run_rule(ea.rules[0], END, START)
    mod.process.assert_called_with({'@timestamp': END})
    assert add_alert.call_count == 1

    # Assert that dropmatchexception behaves properly
    mod.process = mock.MagicMock(side_effect=DropMatchException)
    ea.rules[0]['type'].matches = [{'@timestamp': END}]
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        with mock.patch.object(ea, 'add_aggregated_alert') as add_alert:
            ea.run_rule(ea.rules[0], END, START)
    mod.process.assert_called_with({'@timestamp': END})
    assert add_alert.call_count == 0


def test_agg(ea):
    ea.max_aggregation = 1337
    hits_timestamps = ['2014-09-26T12:34:45', '2014-09-26T12:40:45', '2014-09-26T12:47:45']
    alerttime1 = dt_to_ts(ts_to_dt(hits_timestamps[0]) + datetime.timedelta(minutes=10))
    hits = generate_hits(hits_timestamps)
    ea.current_es.search.return_value = hits
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        # Aggregate first two, query over full range
        ea.rules[0]['aggregation'] = datetime.timedelta(minutes=10)
        ea.rules[0]['type'].matches = [{'@timestamp': h} for h in hits_timestamps]
        ea.run_rule(ea.rules[0], END, START)

    # Assert that the three matches were added to elasticsearch
    call1 = ea.writeback_es.create.call_args_list[0][1]['body']
    call2 = ea.writeback_es.create.call_args_list[1][1]['body']
    call3 = ea.writeback_es.create.call_args_list[2][1]['body']
    assert call1['match_body'] == {'@timestamp': '2014-09-26T12:34:45'}
    assert not call1['alert_sent']
    assert 'aggregate_id' not in call1
    assert call1['alert_time'] == alerttime1

    assert call2['match_body'] == {'@timestamp': '2014-09-26T12:40:45'}
    assert not call2['alert_sent']
    assert call2['aggregate_id'] == 'ABCD'

    assert call3['match_body'] == {'@timestamp': '2014-09-26T12:47:45'}
    assert not call3['alert_sent']
    assert 'aggregate_id' not in call3

    # First call - Find all pending alerts (only entries without agg_id)
    # Second call - Find matches with agg_id == 'ABCD'
    # Third call - Find matches with agg_id == 'CDEF'
    ea.writeback_es.search.side_effect = [{'hits': {'hits': [{'_id': 'ABCD', '_source': call1},
                                                             {'_id': 'CDEF', '_source': call3}]}},
                                          {'hits': {'hits': [{'_id': 'BCDE', '_source': call2}]}},
                                          {'hits': {'total': 0, 'hits': []}}]

    with mock.patch('elastalert.elastalert.elasticsearch_client') as mock_es:
        ea.send_pending_alerts()
        # Assert that current_es was refreshed from the aggregate rules
        assert mock_es.called_with(host='', port='')
        assert mock_es.call_count == 2
    assert_alerts(ea, [hits_timestamps[:2], hits_timestamps[2:]])

    call1 = ea.writeback_es.search.call_args_list[4][1]['body']
    call2 = ea.writeback_es.search.call_args_list[5][1]['body']
    call3 = ea.writeback_es.search.call_args_list[6][1]['body']
    call4 = ea.writeback_es.search.call_args_list[7][1]['body']

    assert 'alert_time' in call2['filter']['range']
    assert call3['query']['query_string']['query'] == 'aggregate_id:ABCD'
    assert call4['query']['query_string']['query'] == 'aggregate_id:CDEF'
    assert ea.writeback_es.search.call_args_list[7][1]['size'] == 1337


def test_agg_cron(ea):
    ea.max_aggregation = 1337
    hits_timestamps = ['2014-09-26T12:34:45', '2014-09-26T12:40:45', '2014-09-26T12:47:45']
    hits = generate_hits(hits_timestamps)
    ea.current_es.search.return_value = hits
    alerttime1 = dt_to_ts(ts_to_dt('2014-09-26T12:46:00'))
    alerttime2 = dt_to_ts(ts_to_dt('2014-09-26T13:04:00'))

    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        with mock.patch('elastalert.elastalert.croniter.get_next') as mock_ts:
            # Aggregate first two, query over full range
            mock_ts.side_effect = [dt_to_unix(ts_to_dt('2014-09-26T12:46:00')), dt_to_unix(ts_to_dt('2014-09-26T13:04:00'))]
            ea.rules[0]['aggregation'] = {'schedule': '*/5 * * * *'}
            ea.rules[0]['type'].matches = [{'@timestamp': h} for h in hits_timestamps]
            ea.run_rule(ea.rules[0], END, START)

    # Assert that the three matches were added to elasticsearch
    call1 = ea.writeback_es.create.call_args_list[0][1]['body']
    call2 = ea.writeback_es.create.call_args_list[1][1]['body']
    call3 = ea.writeback_es.create.call_args_list[2][1]['body']

    assert call1['match_body'] == {'@timestamp': '2014-09-26T12:34:45'}
    assert not call1['alert_sent']
    assert 'aggregate_id' not in call1
    assert call1['alert_time'] == alerttime1

    assert call2['match_body'] == {'@timestamp': '2014-09-26T12:40:45'}
    assert not call2['alert_sent']
    assert call2['aggregate_id'] == 'ABCD'

    assert call3['match_body'] == {'@timestamp': '2014-09-26T12:47:45'}
    assert call3['alert_time'] == alerttime2
    assert not call3['alert_sent']
    assert 'aggregate_id' not in call3


def test_agg_no_writeback_connectivity(ea):
    """ Tests that if writeback_es throws an exception, the matches will be added to 'agg_matches' and when
    run again, that they will be passed again to add_aggregated_alert """
    hit1, hit2, hit3 = '2014-09-26T12:34:45', '2014-09-26T12:40:45', '2014-09-26T12:47:45'
    hits = generate_hits([hit1, hit2, hit3])
    ea.current_es.search.return_value = hits
    ea.rules[0]['aggregation'] = datetime.timedelta(minutes=10)
    ea.rules[0]['type'].matches = [{'@timestamp': hit1},
                                   {'@timestamp': hit2},
                                   {'@timestamp': hit3}]
    ea.writeback_es.create.side_effect = elasticsearch.exceptions.ElasticsearchException('Nope')
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        with mock.patch.object(ea, 'find_pending_aggregate_alert', return_value=None):
            ea.run_rule(ea.rules[0], END, START)

    assert ea.rules[0]['agg_matches'] == [{'@timestamp': hit1},
                                          {'@timestamp': hit2},
                                          {'@timestamp': hit3}]

    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.add_aggregated_alert = mock.Mock()

    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        ea.run_rule(ea.rules[0], END, START)

    ea.add_aggregated_alert.assert_any_call({'@timestamp': hit1}, ea.rules[0])
    ea.add_aggregated_alert.assert_any_call({'@timestamp': hit2}, ea.rules[0])
    ea.add_aggregated_alert.assert_any_call({'@timestamp': hit3}, ea.rules[0])


def test_silence(ea):<commit_msg>Add a test for usage of query_key in an aggregate rule<commit_after># -*- coding: utf-8 -*-
import contextlib
import copy
import datetime
import json
import threading

import elasticsearch
import mock
import pytest
from elasticsearch.exceptions import ElasticsearchException

from elastalert.enhancements import BaseEnhancement
from elastalert.enhancements import DropMatchException
from elastalert.kibana import dashboard_temp
from elastalert.util import dt_to_ts
from elastalert.util import dt_to_unix
from elastalert.util import dt_to_unixms
from elastalert.util import EAException
from elastalert.util import ts_to_dt
from elastalert.util import unix_to_dt


START_TIMESTAMP = '2014-09-26T12:34:45Z'
END_TIMESTAMP = '2014-09-27T12:34:45Z'
START = ts_to_dt(START_TIMESTAMP)
END = ts_to_dt(END_TIMESTAMP)


def _set_hits(ea_inst, hits):
    res = {'hits': {'total': len(hits), 'hits': hits}}
    ea_inst.client_es.return_value = res


def generate_hits(timestamps, **kwargs):
    hits = []
    id_iter = xrange(len(timestamps)).__iter__()
    for ts in timestamps:
        data = {'_id': 'id' + str(id_iter.next()),
                '_source': {'@timestamp': ts},
                '_type': 'logs',
                '_index': 'idx'}
        for key, item in kwargs.iteritems():
            data['_source'][key] = item
        # emulate process_hits(), add metadata to _source
        for field in ['_id', '_type', '_index']:
            data['_source'][field] = data[field]
        hits.append(data)
    return {'hits': {'total': len(hits), 'hits': hits}}


def assert_alerts(ea_inst, calls):
    """ Takes a list of lists of timestamps. Asserts that an alert was called for each list, containing those timestamps. """
    assert ea_inst.rules[0]['alert'][0].alert.call_count == len(calls)
    for call_num, call_args in enumerate(ea_inst.rules[0]['alert'][0].alert.call_args_list):
        assert not any([match['@timestamp'] not in calls[call_num] for match in call_args[0][0]])
        assert len(call_args[0][0]) == len(calls[call_num])


def test_starttime(ea):
    invalid = ['2014-13-13',
               '2014-11-24T30:00:00',
               'Not A Timestamp']
    for ts in invalid:
        with pytest.raises((TypeError, ValueError)):
            ts_to_dt(ts)


def test_init_rule(ea):
    # Simulate state of a rule just loaded from a file
    ea.rules[0]['minimum_starttime'] = datetime.datetime.now()
    new_rule = copy.copy(ea.rules[0])
    map(new_rule.pop, ['agg_matches', 'current_aggregate_id', 'processed_hits', 'minimum_starttime'])

    # Properties are copied from ea.rules[0]
    ea.rules[0]['starttime'] = '2014-01-02T00:11:22'
    ea.rules[0]['processed_hits'] = ['abcdefg']
    new_rule = ea.init_rule(new_rule, False)
    for prop in ['starttime', 'agg_matches', 'current_aggregate_id', 'processed_hits', 'minimum_starttime']:
        assert new_rule[prop] == ea.rules[0][prop]

    # Properties are fresh
    new_rule = ea.init_rule(new_rule, True)
    new_rule.pop('starttime')
    assert 'starttime' not in new_rule
    assert new_rule['processed_hits'] == {}


def test_query(ea):
    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.run_query(ea.rules[0], START, END)
    ea.current_es.search.assert_called_with(body={'query': {'filtered': {'filter': {'bool': {'must': [{'range': {'@timestamp': {'lte': END_TIMESTAMP, 'gt': START_TIMESTAMP}}}]}}}}, 'sort': [{'@timestamp': {'order': 'asc'}}]}, index='idx', _source_include=['@timestamp'], ignore_unavailable=True, size=ea.rules[0]['max_query_size'], scroll=ea.conf['scroll_keepalive'])


def test_query_with_fields(ea):
    ea.rules[0]['_source_enabled'] = False
    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.run_query(ea.rules[0], START, END)
    ea.current_es.search.assert_called_with(body={'query': {'filtered': {'filter': {'bool': {'must': [{'range': {'@timestamp': {'lte': END_TIMESTAMP, 'gt': START_TIMESTAMP}}}]}}}}, 'sort': [{'@timestamp': {'order': 'asc'}}], 'fields': ['@timestamp']}, index='idx', ignore_unavailable=True, size=ea.rules[0]['max_query_size'], scroll=ea.conf['scroll_keepalive'])


def test_query_with_unix(ea):
    ea.rules[0]['timestamp_type'] = 'unix'
    ea.rules[0]['dt_to_ts'] = dt_to_unix
    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.run_query(ea.rules[0], START, END)
    start_unix = dt_to_unix(START)
    end_unix = dt_to_unix(END)
    ea.current_es.search.assert_called_with(body={'query': {'filtered': {'filter': {'bool': {'must': [{'range': {'@timestamp': {'lte': end_unix, 'gt': start_unix}}}]}}}}, 'sort': [{'@timestamp': {'order': 'asc'}}]}, index='idx', _source_include=['@timestamp'], ignore_unavailable=True, size=ea.rules[0]['max_query_size'], scroll=ea.conf['scroll_keepalive'])


def test_query_with_unixms(ea):
    ea.rules[0]['timestamp_type'] = 'unixms'
    ea.rules[0]['dt_to_ts'] = dt_to_unixms
    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.run_query(ea.rules[0], START, END)
    start_unix = dt_to_unixms(START)
    end_unix = dt_to_unixms(END)
    ea.current_es.search.assert_called_with(body={'query': {'filtered': {'filter': {'bool': {'must': [{'range': {'@timestamp': {'lte': end_unix, 'gt': start_unix}}}]}}}}, 'sort': [{'@timestamp': {'order': 'asc'}}]}, index='idx', _source_include=['@timestamp'], ignore_unavailable=True, size=ea.rules[0]['max_query_size'], scroll=ea.conf['scroll_keepalive'])


def test_no_hits(ea):
    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.run_query(ea.rules[0], START, END)
    assert ea.rules[0]['type'].add_data.call_count == 0


def test_no_terms_hits(ea):
    ea.rules[0]['use_terms_query'] = True
    ea.rules[0]['query_key'] = 'QWERTY'
    ea.rules[0]['doc_type'] = 'uiop'
    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.run_query(ea.rules[0], START, END)
    assert ea.rules[0]['type'].add_terms_data.call_count == 0


def test_some_hits(ea):
    hits = generate_hits([START_TIMESTAMP, END_TIMESTAMP])
    hits_dt = generate_hits([START, END])
    ea.current_es.search.return_value = hits
    ea.run_query(ea.rules[0], START, END)
    assert ea.rules[0]['type'].add_data.call_count == 1
    ea.rules[0]['type'].add_data.assert_called_with([x['_source'] for x in hits_dt['hits']['hits']])


def test_some_hits_unix(ea):
    ea.rules[0]['timestamp_type'] = 'unix'
    ea.rules[0]['dt_to_ts'] = dt_to_unix
    ea.rules[0]['ts_to_dt'] = unix_to_dt
    hits = generate_hits([dt_to_unix(START), dt_to_unix(END)])
    hits_dt = generate_hits([START, END])
    ea.current_es.search.return_value = copy.deepcopy(hits)
    ea.run_query(ea.rules[0], START, END)
    assert ea.rules[0]['type'].add_data.call_count == 1
    ea.rules[0]['type'].add_data.assert_called_with([x['_source'] for x in hits_dt['hits']['hits']])


def _duplicate_hits_generator(timestamps, **kwargs):
    """Generator repeatedly returns identical hits dictionaries
    """
    while True:
        yield generate_hits(timestamps, **kwargs)


def test_duplicate_timestamps(ea):
    ea.current_es.search.side_effect = _duplicate_hits_generator([START_TIMESTAMP] * 3, blah='duplicate')
    ea.run_query(ea.rules[0], START, ts_to_dt('2014-01-01T00:00:00Z'))

    assert len(ea.rules[0]['type'].add_data.call_args_list[0][0][0]) == 3
    assert ea.rules[0]['type'].add_data.call_count == 1

    # Run the query again, duplicates will be removed and not added
    ea.run_query(ea.rules[0], ts_to_dt('2014-01-01T00:00:00Z'), END)
    assert ea.rules[0]['type'].add_data.call_count == 1


def test_match(ea):
    hits = generate_hits([START_TIMESTAMP, END_TIMESTAMP])
    ea.current_es.search.return_value = hits
    ea.rules[0]['type'].matches = [{'@timestamp': END}]
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        ea.run_rule(ea.rules[0], END, START)

    ea.rules[0]['alert'][0].alert.called_with({'@timestamp': END_TIMESTAMP})
    assert ea.rules[0]['alert'][0].alert.call_count == 1


def test_run_rule_calls_garbage_collect(ea):
    start_time = '2014-09-26T00:00:00Z'
    end_time = '2014-09-26T12:00:00Z'
    ea.buffer_time = datetime.timedelta(hours=1)
    ea.run_every = datetime.timedelta(hours=1)
    with contextlib.nested(mock.patch.object(ea.rules[0]['type'], 'garbage_collect'),
                           mock.patch.object(ea, 'run_query')) as (mock_gc, mock_get_hits):
        ea.run_rule(ea.rules[0], ts_to_dt(end_time), ts_to_dt(start_time))

    # Running elastalert every hour for 12 hours, we should see self.garbage_collect called 12 times.
    assert mock_gc.call_count == 12

    # The calls should be spaced 1 hour apart
    expected_calls = [ts_to_dt(start_time) + datetime.timedelta(hours=i) for i in range(1, 13)]
    for e in expected_calls:
        mock_gc.assert_any_call(e)


def run_rule_query_exception(ea, mock_es):
    with mock.patch('elastalert.elastalert.elasticsearch_client') as mock_es_init:
        mock_es_init.return_value = mock_es
        ea.run_rule(ea.rules[0], END, START)

    # Assert neither add_data nor garbage_collect were called
    # and that starttime did not change
    assert ea.rules[0].get('starttime') == START
    assert ea.rules[0]['type'].add_data.call_count == 0
    assert ea.rules[0]['type'].garbage_collect.call_count == 0
    assert ea.rules[0]['type'].add_count_data.call_count == 0


def test_query_exception(ea):
    mock_es = mock.Mock()
    mock_es.search.side_effect = ElasticsearchException
    run_rule_query_exception(ea, mock_es)


def test_query_exception_count_query(ea):
    ea.rules[0]['use_count_query'] = True
    ea.rules[0]['doc_type'] = 'blahblahblahblah'
    mock_es = mock.Mock()
    mock_es.count.side_effect = ElasticsearchException
    run_rule_query_exception(ea, mock_es)


def test_match_with_module(ea):
    mod = BaseEnhancement(ea.rules[0])
    mod.process = mock.Mock()
    ea.rules[0]['match_enhancements'] = [mod]
    test_match(ea)
    mod.process.assert_called_with({'@timestamp': END})


def test_match_with_module_with_agg(ea):
    mod = BaseEnhancement(ea.rules[0])
    mod.process = mock.Mock()
    ea.rules[0]['match_enhancements'] = [mod]
    ea.rules[0]['aggregation'] = datetime.timedelta(minutes=15)
    hits = generate_hits([START_TIMESTAMP, END_TIMESTAMP])
    ea.current_es.search.return_value = hits
    ea.rules[0]['type'].matches = [{'@timestamp': END}]
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        ea.run_rule(ea.rules[0], END, START)
    assert mod.process.call_count == 0


def test_match_with_enhancements_first(ea):
    mod = BaseEnhancement(ea.rules[0])
    mod.process = mock.Mock()
    ea.rules[0]['match_enhancements'] = [mod]
    ea.rules[0]['aggregation'] = datetime.timedelta(minutes=15)
    ea.rules[0]['run_enhancements_first'] = True
    hits = generate_hits([START_TIMESTAMP, END_TIMESTAMP])
    ea.current_es.search.return_value = hits
    ea.rules[0]['type'].matches = [{'@timestamp': END}]
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        with mock.patch.object(ea, 'add_aggregated_alert') as add_alert:
            ea.run_rule(ea.rules[0], END, START)
    mod.process.assert_called_with({'@timestamp': END})
    assert add_alert.call_count == 1

    # Assert that dropmatchexception behaves properly
    mod.process = mock.MagicMock(side_effect=DropMatchException)
    ea.rules[0]['type'].matches = [{'@timestamp': END}]
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        with mock.patch.object(ea, 'add_aggregated_alert') as add_alert:
            ea.run_rule(ea.rules[0], END, START)
    mod.process.assert_called_with({'@timestamp': END})
    assert add_alert.call_count == 0


def test_agg(ea):
    ea.max_aggregation = 1337
    hits_timestamps = ['2014-09-26T12:34:45', '2014-09-26T12:40:45', '2014-09-26T12:47:45']
    alerttime1 = dt_to_ts(ts_to_dt(hits_timestamps[0]) + datetime.timedelta(minutes=10))
    hits = generate_hits(hits_timestamps)
    ea.current_es.search.return_value = hits
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        # Aggregate first two, query over full range
        ea.rules[0]['aggregation'] = datetime.timedelta(minutes=10)
        ea.rules[0]['type'].matches = [{'@timestamp': h} for h in hits_timestamps]
        ea.run_rule(ea.rules[0], END, START)

    # Assert that the three matches were added to elasticsearch
    call1 = ea.writeback_es.create.call_args_list[0][1]['body']
    call2 = ea.writeback_es.create.call_args_list[1][1]['body']
    call3 = ea.writeback_es.create.call_args_list[2][1]['body']
    assert call1['match_body'] == {'@timestamp': '2014-09-26T12:34:45'}
    assert not call1['alert_sent']
    assert 'aggregate_id' not in call1
    assert call1['alert_time'] == alerttime1

    assert call2['match_body'] == {'@timestamp': '2014-09-26T12:40:45'}
    assert not call2['alert_sent']
    assert call2['aggregate_id'] == 'ABCD'

    assert call3['match_body'] == {'@timestamp': '2014-09-26T12:47:45'}
    assert not call3['alert_sent']
    assert 'aggregate_id' not in call3

    # First call - Find all pending alerts (only entries without agg_id)
    # Second call - Find matches with agg_id == 'ABCD'
    # Third call - Find matches with agg_id == 'CDEF'
    ea.writeback_es.search.side_effect = [{'hits': {'hits': [{'_id': 'ABCD', '_source': call1},
                                                             {'_id': 'CDEF', '_source': call3}]}},
                                          {'hits': {'hits': [{'_id': 'BCDE', '_source': call2}]}},
                                          {'hits': {'total': 0, 'hits': []}}]

    with mock.patch('elastalert.elastalert.elasticsearch_client') as mock_es:
        ea.send_pending_alerts()
        # Assert that current_es was refreshed from the aggregate rules
        assert mock_es.called_with(host='', port='')
        assert mock_es.call_count == 2
    assert_alerts(ea, [hits_timestamps[:2], hits_timestamps[2:]])

    call1 = ea.writeback_es.search.call_args_list[4][1]['body']
    call2 = ea.writeback_es.search.call_args_list[5][1]['body']
    call3 = ea.writeback_es.search.call_args_list[6][1]['body']
    call4 = ea.writeback_es.search.call_args_list[7][1]['body']

    assert 'alert_time' in call2['filter']['range']
    assert call3['query']['query_string']['query'] == 'aggregate_id:ABCD'
    assert call4['query']['query_string']['query'] == 'aggregate_id:CDEF'
    assert ea.writeback_es.search.call_args_list[7][1]['size'] == 1337


def test_agg_cron(ea):
    ea.max_aggregation = 1337
    hits_timestamps = ['2014-09-26T12:34:45', '2014-09-26T12:40:45', '2014-09-26T12:47:45']
    hits = generate_hits(hits_timestamps)
    ea.current_es.search.return_value = hits
    alerttime1 = dt_to_ts(ts_to_dt('2014-09-26T12:46:00'))
    alerttime2 = dt_to_ts(ts_to_dt('2014-09-26T13:04:00'))

    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        with mock.patch('elastalert.elastalert.croniter.get_next') as mock_ts:
            # Aggregate first two, query over full range
            mock_ts.side_effect = [dt_to_unix(ts_to_dt('2014-09-26T12:46:00')), dt_to_unix(ts_to_dt('2014-09-26T13:04:00'))]
            ea.rules[0]['aggregation'] = {'schedule': '*/5 * * * *'}
            ea.rules[0]['type'].matches = [{'@timestamp': h} for h in hits_timestamps]
            ea.run_rule(ea.rules[0], END, START)

    # Assert that the three matches were added to elasticsearch
    call1 = ea.writeback_es.create.call_args_list[0][1]['body']
    call2 = ea.writeback_es.create.call_args_list[1][1]['body']
    call3 = ea.writeback_es.create.call_args_list[2][1]['body']

    assert call1['match_body'] == {'@timestamp': '2014-09-26T12:34:45'}
    assert not call1['alert_sent']
    assert 'aggregate_id' not in call1
    assert call1['alert_time'] == alerttime1

    assert call2['match_body'] == {'@timestamp': '2014-09-26T12:40:45'}
    assert not call2['alert_sent']
    assert call2['aggregate_id'] == 'ABCD'

    assert call3['match_body'] == {'@timestamp': '2014-09-26T12:47:45'}
    assert call3['alert_time'] == alerttime2
    assert not call3['alert_sent']
    assert 'aggregate_id' not in call3


def test_agg_no_writeback_connectivity(ea):
    """ Tests that if writeback_es throws an exception, the matches will be added to 'agg_matches' and when
    run again, that they will be passed again to add_aggregated_alert """
    hit1, hit2, hit3 = '2014-09-26T12:34:45', '2014-09-26T12:40:45', '2014-09-26T12:47:45'
    hits = generate_hits([hit1, hit2, hit3])
    ea.current_es.search.return_value = hits
    ea.rules[0]['aggregation'] = datetime.timedelta(minutes=10)
    ea.rules[0]['type'].matches = [{'@timestamp': hit1},
                                   {'@timestamp': hit2},
                                   {'@timestamp': hit3}]
    ea.writeback_es.create.side_effect = elasticsearch.exceptions.ElasticsearchException('Nope')
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        with mock.patch.object(ea, 'find_pending_aggregate_alert', return_value=None):
            ea.run_rule(ea.rules[0], END, START)

    assert ea.rules[0]['agg_matches'] == [{'@timestamp': hit1},
                                          {'@timestamp': hit2},
                                          {'@timestamp': hit3}]

    ea.current_es.search.return_value = {'hits': {'total': 0, 'hits': []}}
    ea.add_aggregated_alert = mock.Mock()

    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        ea.run_rule(ea.rules[0], END, START)

    ea.add_aggregated_alert.assert_any_call({'@timestamp': hit1}, ea.rules[0])
    ea.add_aggregated_alert.assert_any_call({'@timestamp': hit2}, ea.rules[0])
    ea.add_aggregated_alert.assert_any_call({'@timestamp': hit3}, ea.rules[0])


def test_agg_with_query_key(ea):
    ea.max_aggregation = 1337
    hits_timestamps = ['2014-09-26T12:34:45', '2014-09-26T12:40:45', '2014-09-26T12:43:45']
    alerttime1 = dt_to_ts(ts_to_dt(hits_timestamps[0]) + datetime.timedelta(minutes=10))
    alerttime2 = dt_to_ts(ts_to_dt(hits_timestamps[1]) + datetime.timedelta(minutes=10))
    hits = generate_hits(hits_timestamps)
    ea.current_es.search.return_value = hits
    with mock.patch('elastalert.elastalert.elasticsearch_client'):
        # Aggregate first two, query over full range
        ea.rules[0]['aggregation'] = datetime.timedelta(minutes=10)
        ea.rules[0]['type'].matches = [{'@timestamp': h} for h in hits_timestamps]
        ea.rules[0]['type'].matches[0]['key'] = 'Key Value 1'
        ea.rules[0]['type'].matches[1]['key'] = 'Key Value 2'
        ea.rules[0]['type'].matches[2]['key'] = 'Key Value 1'
        ea.rules[0]['query_key'] = 'key'
        ea.run_rule(ea.rules[0], END, START)

    # Assert that the three matches were added to elasticsearch
    call1 = ea.writeback_es.create.call_args_list[0][1]['body']
    call2 = ea.writeback_es.create.call_args_list[1][1]['body']
    call3 = ea.writeback_es.create.call_args_list[2][1]['body']
    assert call1['match_body'] == {'@timestamp': '2014-09-26T12:34:45', 'key': 'Key Value 1'}
    assert not call1['alert_sent']
    assert 'aggregate_id' not in call1
    assert 'aggregate_key' in call1
    assert call1['aggregate_key'] == 'Key Value 1'
    assert call1['alert_time'] == alerttime1

    assert call2['match_body'] == {'@timestamp': '2014-09-26T12:40:45', 'key': 'Key Value 2'}
    assert not call2['alert_sent']
    assert 'aggregate_id' not in call2
    assert 'aggregate_key' in call2
    assert call2['aggregate_key'] == 'Key Value 2'
    assert call2['alert_time'] == alerttime2

    assert call3['match_body'] == {'@timestamp': '2014-09-26T12:43:45', 'key': 'Key Value 1', 'key': 'Key Value 1'}
    assert not call3['alert_sent']
    assert call3['aggregate_id'] == 'ABCD'
    assert 'aggregate_key' in call3
    assert call3['aggregate_key'] == 'Key Value 1'
    assert call3['alert_time'] == alerttime1

    # First call - Find all pending alerts (only entries without agg_id)
    # Second call - Find matches with agg_id == 'ABCD'
    # Third call - Find matches with agg_id == 'CDEF'
    ea.writeback_es.search.side_effect = [{'hits': {'hits': [{'_id': 'ABCD', '_source': call1},
                                                             {'_id': 'CDEF', '_source': call3}]}},
                                          {'hits': {'hits': [{'_id': 'BCDE', '_source': call2}]}},
                                          {'hits': {'total': 0, 'hits': []}}]

    with mock.patch('elastalert.elastalert.elasticsearch_client') as mock_es:
        ea.send_pending_alerts()
        # Assert that current_es was refreshed from the aggregate rules
        assert mock_es.called_with(host='', port='')
        assert mock_es.call_count == 2
    assert_alerts(ea, [hits_timestamps[:2], hits_timestamps[2:]])

    call1 = ea.writeback_es.search.call_args_list[4][1]['body']
    call2 = ea.writeback_es.search.call_args_list[5][1]['body']
    call3 = ea.writeback_es.search.call_args_list[6][1]['body']
    call4 = ea.writeback_es.search.call_args_list[7][1]['body']

    assert 'alert_time' in call2['filter']['range']
    assert call3['query']['query_string']['query'] == 'aggregate_id:ABCD'
    assert call4['query']['query_string']['query'] == 'aggregate_id:CDEF'
    assert ea.writeback_es.search.call_args_list[7][1]['size'] == 1337


def test_silence(ea):
{'fraction_non_alphanumeric': 0.16307418246596259, 'fraction_numerical': 0.036238707214658356, 'mean_word_length': 4.1007268951194185, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 30, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '307228', 'n_tokens_mistral': 14932, 'n_tokens_neox': 13886, 'n_words': 2472}