import unittest

from fgai.graylog_source import GraylogStreamSource


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

    def call_tool(self, _name, _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):
        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")


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