63 lines
2.5 KiB
Python
63 lines
2.5 KiB
Python
import unittest
|
|
|
|
from fgai.llm import _json_object_from_text, compact_dashboard_analysis
|
|
|
|
|
|
class LlmTests(unittest.TestCase):
|
|
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"]), 20)
|
|
self.assertEqual(compact["field_deviations"][0]["entity"], "alice")
|
|
self.assertNotIn("sample_events", compact["field_deviations"][0])
|
|
self.assertNotIn("graylog_query", str(compact))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|