Fix test_modules to add more time with better event waiting (#51312)

Author: blakerouse

filebeat/tests/system/test_modules.py

@@ -5,13 +5,20 @@
 import unittest
 import glob
 import subprocess
+import time
 
 import json
 import logging
 from parameterized import parameterized
 from elasticsearch import Elasticsearch, NotFoundError
 from deepdiff import DeepDiff
 
+# Maximum time (seconds) Filebeat is allowed to ingest a module's test file.
+# Inputs started with --once exit as soon as the file is fully read; tailing
+# inputs (journald, filestream) never exit, so the test stops Filebeat once the
+# expected events have been indexed or this timeout elapses.
+MODULE_INGEST_TIMEOUT = 60
+
 # datasets for which @timestamp is removed due to date missing
 remove_timestamp = {
     "activemq.audit",
@@ -258,17 +265,19 @@ def run_on_file(self, module, fileset, test_file, cfgfile):
                                     stdout=output,
                                     stderr=subprocess.STDOUT,
                                     bufsize=0)
-            # The journald input (used by some modules like 'system') does not
-            # support the --once flag, hence we run Filebeat for at most
-            # 15 seconds, if it does not finish, then kill the process.
-            # If for any reason the Filebeat process gets stuck, only SIGKILL
-            # will terminate it. We use SIGKILL to avoid leaking any running
-            # process that could interfere with other tests
-            try:
-                proc.wait(15)
-            except subprocess.TimeoutExpired:
-                # Send SIGKILL
-                proc.kill()
+            if "--once" in cmd:
+                # Process will exit on its own once the file is fully read. Wait the maximum time,
+                # if it doesn't exit by then, kill it to avoid leaking a process.
+                try:
+                    proc.wait(MODULE_INGEST_TIMEOUT)
+                except subprocess.TimeoutExpired:
+                    proc.kill()
+                    proc.wait()
+            else:
+                # The journald and filestream inputs do not support --once and
+                # tail the file forever. Stop filebeat once elasticsearch has
+                # indexed the expected number of events (or the timeout elapses).
+                self._wait_for_events_then_stop(proc, test_file, MODULE_INGEST_TIMEOUT)
 
         # List of errors to check in filebeat output logs
         errors = ["error loading pipeline for fileset"]
@@ -311,6 +320,38 @@ def run_on_file(self, module, fileset, test_file, cfgfile):
 
         self._test_expected_events(test_file, objects)
 
+    def _wait_for_events_then_stop(self, proc, test_file, timeout):
+        expected = self._expected_event_count(test_file)
+        deadline = time.monotonic() + timeout
+        try:
+            while time.monotonic() < deadline:
+                if proc.poll() is not None:
+                    # Filebeat exited on its own; nothing left to stop.
+                    return
+                if expected is not None and self._indexed_event_count() >= expected:
+                    break
+                time.sleep(0.5)
+        finally:
+            if proc.poll() is None:
+                proc.kill()
+            proc.wait()
+
+    def _expected_event_count(self, test_file):
+        if os.getenv("GENERATE"):
+            return None
+        try:
+            with open(test_file + "-expected.json", "r") as f:
+                return len(json.load(f))
+        except (FileNotFoundError, ValueError):
+            return None
+
+    def _indexed_event_count(self):
+        try:
+            self.es.indices.refresh(index=self.index_name)
+            return self.es.count(index=self.index_name)["count"]
+        except NotFoundError:
+            return 0
+
     def _test_expected_events(self, test_file, objects):
 
         # Generate expected files if GENERATE env variable is set