improve agregate with status

tests/test_graylog_aggregate.py

@@ -4,12 +4,16 @@ from fgai.graylog_aggregate import GraylogAggregateSource
 
 
 class _AggregateClient:
-    def __init__(self, responses):
+    def __init__(self, responses, schema=None):
         self.responses = list(responses)
         self.arguments = []
+        self.schema = schema
 
     def probe(self):
-        return {"status": "connected"}
+        status = {"status": "connected"}
+        if self.schema:
+            status["tool_schemas"] = {"aggregate_messages": self.schema}
+        return status
 
     def call_tool(self, _name, arguments):
         self.arguments.append(arguments)
@@ -41,6 +45,18 @@ class GraylogAggregateTests(unittest.TestCase):
         self.assertEqual(status["aggregate_events"], 42)
         self.assertEqual(len(client.arguments), 2)
 
+    def test_uses_tool_schema_to_avoid_unsupported_fields(self):
+        schema = {"properties": {"query": {}, "streams": {}, "range_seconds": {}, "series": {}}}
+        client = _AggregateClient([
+            {"result": {"content": [{"type": "text", "text": '{"schema":[{"name":"count"}],"datarows":[[7]]}'}]}}
+        ], schema=schema)
+
+        status = GraylogAggregateSource(client, "firewall").fetch_count()
+
+        self.assertEqual(status["aggregate_events"], 7)
+        self.assertIn("series", client.arguments[0])
+        self.assertNotIn("group_by", client.arguments[0])
+
 
 if __name__ == "__main__":
     unittest.main()