From d83120817ee9eb32bb7fd2ccd8af6bc5fa3deca7 Mon Sep 17 00:00:00 2001 From: Emma Rose Date: Tue, 30 Jul 2019 17:39:49 -0400 Subject: [PATCH] Wrote tests to interface with elasticsearch --- alerts/geomodel/event.py | 4 +- alerts/geomodel/locality.py | 12 ++++-- alerts/geomodel/query.py | 4 +- tests/alerts/geomodel/test_event.py | 2 +- tests/alerts/geomodel/test_locality.py | 56 ++++++++++++++++++++++---- 5 files changed, 64 insertions(+), 14 deletions(-) diff --git a/alerts/geomodel/event.py b/alerts/geomodel/event.py index d22c0e68..df07591f 100644 --- a/alerts/geomodel/event.py +++ b/alerts/geomodel/event.py @@ -38,7 +38,9 @@ def find_all( for cfg in evt_cfg.queries: search = SearchQuery(minutes=evt_cfg.search_window.minutes) - search.add_must([QSMatch(cfg.lucene)]) + search.add_must(QSMatch(cfg.lucene)) + + print(f'Constructed search query {search}') search_results = query_es(search, evt_cfg.es_index) diff --git a/alerts/geomodel/locality.py b/alerts/geomodel/locality.py index 7b5016d4..21a9e2c2 100644 --- a/alerts/geomodel/locality.py +++ b/alerts/geomodel/locality.py @@ -9,6 +9,10 @@ import alerts.geomodel.query as query # TODO: Switch to dataclasses when we move to Python3.7+ +def _dict_take(dictionary, keys): + return {key: dictionary[key] for key in keys} + + class Locality(NamedTuple): '''Represents a specific locality. ''' @@ -41,16 +45,18 @@ def find_all( def to_state(result: Dict[str, Any]) -> Optional[State]: try: result['localities'] = [ - Locality(**loc) + Locality(**_dict_take(loc, Locality._fields)) for loc in result['localities'] ] - return State(**result) + return State(**_dict_take(result, State._fields)) except TypeError: return None + except KeyError: + return None search = SearchQuery() - search.add_must([TermMatch('type', 'locality')]) + search.add_must([TermMatch('type_', 'locality')]) results = query_es(search, locality.es_index) diff --git a/alerts/geomodel/query.py b/alerts/geomodel/query.py index 26a39d8c..e3738d74 100644 --- a/alerts/geomodel/query.py +++ b/alerts/geomodel/query.py @@ -18,6 +18,8 @@ def wrap(client: ESClient) -> QueryInterface: ''' def wrapper(query: SearchQuery, esindex: str) -> List[Dict[str, Any]]: - return query.execute(client, indices=[esindex]).get('hits', []) + hits = query.execute(client, indices=[esindex]).get('hits', []) + + return [hit.get('_source', {}) for hit in hits] return wrapper diff --git a/tests/alerts/geomodel/test_event.py b/tests/alerts/geomodel/test_event.py index ba361a65..67f7ada6 100644 --- a/tests/alerts/geomodel/test_event.py +++ b/tests/alerts/geomodel/test_event.py @@ -29,7 +29,7 @@ class TestEventElasticSearch(UnitTestSuite): ] for evt in events: - self.populate_test_object(evt) + self.populate_test_event(evt) self.refresh(self.event_index_name) diff --git a/tests/alerts/geomodel/test_locality.py b/tests/alerts/geomodel/test_locality.py index baaa9433..513f2df2 100644 --- a/tests/alerts/geomodel/test_locality.py +++ b/tests/alerts/geomodel/test_locality.py @@ -3,8 +3,48 @@ import unittest import alerts.geomodel.config as config import alerts.geomodel.locality as locality +import alerts.geomodel.query as query from tests.alerts.geomodel.util import query_interface +from tests.unit_test_suite import UnitTestSuite + + +class TestLocalityElasticSearch(UnitTestSuite): + '''Tests for the `locality` module that interact with ES. + ''' + + def test_simple_query(self): + objs = [ + { + 'type_': 'locality', + 'username': 'tester1', + 'localities': [ + { + 'sourceipv4address': '1.2.3.4', + 'city': 'Toronto', + 'country': 'CA', + 'lastaction': datetime.utcnow(), + 'latitude': 43.6529, + 'longitude': -79.3849, + 'radius': 50 + } + ] + } + ] + + for obj in objs: + self.populate_test_event(obj) + + self.refresh(self.event_index_name) + + query_iface = query.wrap(self.es_client) + loc_cfg = config.Localities(self.event_index_name, 30, 50.0) + + results = locality.find_all(query_iface, loc_cfg) + usernames = [state.username for state in results] + + assert len(results) == 1 + assert usernames == ['tester1'] class TestLocality(unittest.TestCase): @@ -14,14 +54,14 @@ class TestLocality(unittest.TestCase): def test_find_all_retrieves_all_states(self): query_iface = query_interface([ { - 'type_': 'geomodel', + 'type_': 'locality', 'username': 'tester1', 'localities': [ { 'sourceipv4address': '1.2.3.4', 'city': 'Toronto', 'country': 'CA', - 'lastaction': datetime.now(), + 'lastaction': datetime.utcnow(), 'latitude': 43.6529, 'longitude': -79.3849, 'radius': 50 @@ -29,14 +69,14 @@ class TestLocality(unittest.TestCase): ] }, { - 'type_': 'geomodel', + 'type_': 'locality', 'username': 'tester2', 'localities': [ { 'sourceipv4address': '4.3.2.1', 'city': 'San Francisco', 'country': 'USA', - 'lastaction': datetime.now(), + 'lastaction': datetime.utcnow(), 'latitude': 37.773972, 'longitude': -122.431297, 'radius': 50 @@ -59,19 +99,19 @@ class TestLocality(unittest.TestCase): query_iface = query_interface([ # Invalid top-level State { - 'type__': 'geomodel', # Should have only one underscore (_) + 'type__': 'locality', # Should have only one underscore (_) 'username': 'tester', 'localities': [] }, # Valid State { - 'type_': 'geomodel', + 'type_': 'locality', 'username': 'validtester', 'localities': [] }, # Invalid locality data { - 'type_': 'geomodel', + 'type_': 'locality', 'username': 'tester2', 'localities': [ { @@ -79,7 +119,7 @@ class TestLocality(unittest.TestCase): 'sourceipaddress': '1.2.3.4', 'city': 'San Francisco', 'country': 'USA', - 'lastaction': datetime.now(), + 'lastaction': datetime.utcnow(), 'latitude': 37.773972, 'longitude': -122.431297, 'radius': 50