Wrote tests to interface with elasticsearch

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)
 

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)
 

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

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)
 

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