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()