import unittest

from fgai.graylog_source import GraylogStreamSource


class _Client:
    def __init__(self):
        self.arguments = None

    def probe(self):
        return {"status": "connected"}

    def call_tool(self, _name, arguments):
        self.arguments = arguments
        return {"result": {"content": [{"type": "text", "text": '{"messages":[{"client":"10.0.0.1","server":"8.8.8.8","result":"deny"}]}' }]}}


class GraylogSourceTests(unittest.TestCase):
    def test_applies_custom_mapping_to_generic_stream_message(self):
        client = _Client()
        events, status = GraylogStreamSource(
            client, "vpn", field_mapping='{"srcip":"client","dstip":"server","action":"result"}'
        ).fetch()
        self.assertEqual(status["events_fetched"], 1)
        self.assertEqual(events[0].src_ip, "10.0.0.1")
        self.assertEqual(events[0].dst_ip, "8.8.8.8")
        self.assertEqual(events[0].action, "deny")
        self.assertEqual(client.arguments["streams"], ["vpn"])
        self.assertEqual(client.arguments["range_seconds"], 300)


if __name__ == "__main__":
    unittest.main()