47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
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"])
|
|
self.assertEqual(client.arguments[0]["metrics"], ["count()"])
|
|
|
|
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()
|