fgAI/tests/test_llm.py

import unittest

from fgai.llm import _compact_profile_suggestion, _json_object_from_text, compact_dashboard_analysis


class LlmTests(unittest.TestCase):
    def test_compact_profile_suggestion_without_fields_only_keeps_everything(self):
        item = {
            "stream_id": "windows",
            "stream_name": "Windows",
            "events": 10,
            "common_fields": [{"field": "username", "coverage": 1.0}, {"field": "hostname", "coverage": 0.9}],
            "discovery": {"top_fields": [{"field": "username"}, {"field": "hostname"}], "selected_fields": {}, "rejected_fields": []},
            "profile": {"entity_fields": ["username"]},
            "field_shapes": {"username": "categorical_low_cardinality", "hostname": "categorical_low_cardinality"},
            "field_confidence": {"username": 80, "hostname": 40},
        }

        compact = _compact_profile_suggestion(item)

        self.assertNotIn("fields_to_classify", compact)
        self.assertEqual(len(compact["common_fields"]), 2)

    def test_compact_profile_suggestion_with_fields_only_trims_to_subset(self):
        item = {
            "stream_id": "windows",
            "stream_name": "Windows",
            "events": 10,
            "common_fields": [{"field": "username", "coverage": 1.0}, {"field": "hostname", "coverage": 0.9}],
            "discovery": {"top_fields": [{"field": "username"}, {"field": "hostname"}], "selected_fields": {}, "rejected_fields": []},
            "profile": {"entity_fields": ["username"]},
            "field_shapes": {"username": "categorical_low_cardinality", "hostname": "categorical_low_cardinality"},
            "field_confidence": {"username": 80, "hostname": 40},
        }

        compact = _compact_profile_suggestion(item, fields_only=["hostname"])

        self.assertEqual(compact["fields_to_classify"], ["hostname"])
        self.assertEqual([row["field"] for row in compact["common_fields"]], ["hostname"])
        self.assertEqual([row["field"] for row in compact["top_fields"]], ["hostname"])
        self.assertEqual(compact["field_shapes"], {"hostname": "categorical_low_cardinality"})
        self.assertEqual(compact["field_confidence"], {"hostname": 40})

    def test_json_object_from_text_accepts_markdown_wrapped_json(self):
        payload = _json_object_from_text(
            '```json\n{"profiles":[{"stream_id":"windows","entity_fields":["username"]}]}\n```'
        )

        self.assertEqual(payload["profiles"][0]["stream_id"], "windows")

    def test_json_object_from_text_extracts_object_from_extra_text(self):
        payload = _json_object_from_text(
            'Here is the profile:\n{"profiles":[{"stream_id":"firewall"}]}\nDone.'
        )

        self.assertEqual(payload["profiles"][0]["stream_id"], "firewall")

    def test_dashboard_compaction_removes_heavy_evidence_payloads(self):
        status = {
            "summary": {"total": 1000},
            "capabilities": {
                "graylog_mcp": {
                    "status": "connected",
                    "streams": [{"stream_id": f"stream-{index}", "events": index} for index in range(100)],
                    "aggregate_events": 1000,
                    "raw_events_fetched": 50,
                }
            },
            "event_context": {
                "source_profiles": [{"entity": f"10.0.0.{index}"} for index in range(20)],
                "related_activity": [{"entity": f"10.0.0.{index}"} for index in range(40)],
            },
            "field_deviations": {
                "alice": [
                    {
                        "stream_name": "Windows",
                        "detector": "new_relationship",
                        "field": "relationship:username->srcip",
                        "score": 90,
                        "reason": "new srcip value",
                        "sample_events": [{"message": "very large raw event", "graylog_query": "username:alice"}],
                    }
                ]
            },
        }

        compact = compact_dashboard_analysis(status)

        self.assertEqual(compact["capabilities"]["graylog_mcp"]["status"], "connected")
        self.assertNotIn("streams", compact["capabilities"]["graylog_mcp"])
        self.assertEqual(len(compact["event_context"]["source_profiles"]), 10)
        self.assertEqual(len(compact["event_context"]["related_activity"]), 12)
        self.assertEqual(compact["field_deviations"][0]["entity"], "alice")
        self.assertNotIn("sample_events", compact["field_deviations"][0])
        self.assertNotIn("graylog_query", str(compact))

    def test_dashboard_compaction_strips_incident_and_correlation_samples(self):
        status = {
            "incidents": [
                {
                    "entity": "10.0.0.5",
                    "score": 90,
                    "severity": "critical",
                    "correlated_streams": ["Firewall", "Proxy"],
                    "evidence": ["large burst of blocked traffic"],
                    "timeline": [{"message": "raw incident evidence", "graylog_query": "srcip:10.0.0.5"}],
                }
            ],
            "cross_source_correlations": [
                {
                    "entity": "10.0.0.5",
                    "streams": ["Firewall", "Proxy"],
                    "events": 20,
                    "samples": [{"message": "raw correlation evidence", "graylog_query": "srcip:10.0.0.5"}],
                }
            ],
            "anomalies": [
                {
                    "subject": "10.0.0.5",
                    "score": 90,
                    "severity": "critical",
                    "reasons": ["burst"],
                    "evidence": {"timed_events": 100, "sample_events": [{"message": "raw anomaly"}]},
                }
            ],
        }

        compact = compact_dashboard_analysis(status)
        compact_text = str(compact)

        self.assertNotIn("timeline", compact_text)
        self.assertNotIn("samples", compact_text)
        self.assertNotIn("sample_events", compact_text)
        self.assertNotIn("graylog_query", compact_text)
        self.assertNotIn("raw incident evidence", compact_text)
        self.assertEqual(compact["incidents"][0]["entity"], "10.0.0.5")


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