add agregated search

tests/test_config.py

@@ -31,6 +31,13 @@ class ConfigTests(unittest.TestCase):
             public = store.update({"graylog_max_events_per_stream": "25000"})
             self.assertEqual(public["graylog_max_events_per_stream"], 25000)
 
+    def test_graylog_fetch_mode_and_raw_sample_events_are_saved(self):
+        with tempfile.TemporaryDirectory() as directory:
+            store = ConfigStore(str(Path(directory) / "config.json"))
+            public = store.update({"graylog_fetch_mode": "aggregate", "graylog_raw_sample_events": "2500"})
+            self.assertEqual(public["graylog_fetch_mode"], "aggregate")
+            self.assertEqual(public["graylog_raw_sample_events"], 2500)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_graylog_aggregate.py

@@ -0,0 +1,45 @@
+import unittest
+
+from fgai.graylog_aggregate import GraylogAggregateSource
+
+
+class _AggregateClient:
+    def __init__(self, responses):
+        self.responses = list(responses)
+        self.arguments = []
+
+    def probe(self):
+        return {"status": "connected"}
+
+    def call_tool(self, _name, arguments):
+        self.arguments.append(arguments)
+        return self.responses.pop(0)
+
+
+class GraylogAggregateTests(unittest.TestCase):
+    def test_reads_count_from_graylog_schema_rows(self):
+        client = _AggregateClient([
+            {"result": {"content": [{"type": "text", "text": '{"schema":[{"name":"metric: count()"}],"datarows":[[12345]]}'}]}}
+        ])
+
+        status = GraylogAggregateSource(client, "firewall").fetch_count(range_seconds=300)
+
+        self.assertEqual(status["aggregate_status"], "ok")
+        self.assertEqual(status["aggregate_events"], 12345)
+        self.assertEqual(client.arguments[0]["streams"], ["firewall"])
+
+    def test_tries_fallback_argument_shape_after_tool_error(self):
+        client = _AggregateClient([
+            {"result": {"isError": True, "content": [{"type": "text", "text": "bad metrics"}]}},
+            {"result": {"content": [{"type": "text", "text": '{"events": 42}'}]}},
+        ])
+
+        status = GraylogAggregateSource(client, "firewall").fetch_count()
+
+        self.assertEqual(status["aggregate_status"], "ok")
+        self.assertEqual(status["aggregate_events"], 42)
+        self.assertEqual(len(client.arguments), 2)
+
+
+if __name__ == "__main__":
+    unittest.main()