import unittest

from fgai.correlation import correlate_source_ips
from fgai.logs import parse_log_line


class CorrelationTests(unittest.TestCase):
    def test_correlates_same_source_across_streams(self):
        events = [
            parse_log_line("srcip=10.0.0.5 fgai_stream=Fortigate action=blocked"),
            parse_log_line("srcip=10.0.0.5 fgai_stream=DNS action=deny"),
        ]
        result = correlate_source_ips(events)
        self.assertEqual(result[0]["source_ip"], "10.0.0.5")
        self.assertEqual(result[0]["streams"], ["DNS", "Fortigate"])
        self.assertEqual(len(result[0]["samples"]), 2)