Test : filter events

Author: VincentD06

validation/graylog.py

@@ -20,6 +20,15 @@ def _wait(self, condition, attempts, sleep_duration=1):
                 print(self._server.extract_all_logs())
                 raise ServerTimeoutError()
 
+    def _waitWithParam(self, condition, attempts, contion_args, sleep_duration=1):
+        count = 0
+        while not condition(contion_args):
+            time.sleep(sleep_duration)
+            count += 1
+            if count > attempts:
+                print(self._server.extract_all_logs())
+                raise ServerTimeoutError()
+
     def _wait_until_graylog_has_started(self):
         """
         We wait until the default deflector is up, as it seems to be the last operation done on startup
@@ -54,14 +63,25 @@ def create_correlation_count(self, *args, **kwargs):
     def get_events(self):
         return self._api.get_events()
 
-    def get_events_count(self):
-        events = self.get_events()
-        return events['total_events']
-
-    def _has_event(self):
-        events_count = self.get_events_count()
+    def get_events_count(self, event_definition_type=None):
+        response = self.get_events()
+        total = response['total_events']
+        if event_definition_type is None:
+            return total
+        result = 0
+        for i in range(total):
+            event = response['events'][i]['event']
+            if event['event_definition_type'] == event_definition_type:
+                result += 1
+        return result
+
+    def _has_event(self, event_definition_type=None):
+        events_count = self.get_events_count(event_definition_type)
         return events_count == 1
 
-    def wait_until_event(self):
-       self._wait(self._has_event, 60)
+    def wait_until_event(self, event_definition_type=None):
+        if event_definition_type is None:
+            self._wait(self._has_event, 60)
+        else:
+            self._waitWithParam(self._has_event, 60, event_definition_type)
 

validation/test.py

@@ -38,7 +38,7 @@ def test_send_message_should_trigger_correlation_rule(self):
             time.sleep(_PERIOD)
 
             try:
-                self._graylog.wait_until_event()
+                self._graylog.wait_until_event('correlation-count')
             except ServerTimeoutError:
                 print(self._graylog.get_events())
                 events_count = self._graylog.get_events_count()
@@ -52,7 +52,7 @@ def test_send_message_should_trigger_correlation_rule_with_group_by(self):
             time.sleep(_PERIOD)
             inputs.send({'short_message': 'pop'})
 
-            self._graylog.wait_until_event()
+            self._graylog.wait_until_event('correlation-count')
 
     def test_send_message_with_different_values_for_group_by_field_should_not_trigger_correlation_rule_with_group_by(self):
         self._graylog.create_correlation_count(1, group_by=['x'], period=_PERIOD)
@@ -63,7 +63,7 @@ def test_send_message_with_different_values_for_group_by_field_should_not_trigge
             inputs.send({'short_message': 'pop'})
 
             time.sleep(20)
-            self.assertEqual(0, self._graylog.get_events_count())
+            self.assertEqual(0, self._graylog.get_events_count('correlation-count'))
 
     def test_send_message_should_trigger_correlation_rule_with_group_by_when_value_has_a_space__issue27(self):
         self._graylog.create_correlation_count(1, group_by=['x'], period=_PERIOD, messages_order='BEFORE')
@@ -78,7 +78,7 @@ def test_send_message_should_trigger_correlation_rule_with_group_by_when_value_h
             time.sleep(_PERIOD)
             inputs.send({'short_message': 'pop'})
 
-            self._graylog.wait_until_event()
+            self._graylog.wait_until_event('correlation-count')
 
     def test_send_message_should_not_fail_on_correlation_rule_with_group_by_when_value_has_a_double_quote__issue27(self):
         self._graylog.create_correlation_count(1, group_by=['x'], period=_PERIOD, messages_order='BEFORE')