import json
import unittest
from unittest.mock import patch

from fgai.graylog_mcp import GraylogMcpClient


class _Response:
    def __init__(self, body: dict[str, object], session: str = "session-1") -> None:
        self.body = body
        self.headers = {"Mcp-Session-Id": session}

    def read(self) -> bytes:
        return json.dumps(self.body).encode("utf-8")

    def __enter__(self):
        return self

    def __exit__(self, *_args):
        return False


class GraylogMcpTests(unittest.TestCase):
    def test_probe_initializes_and_lists_tools(self):
        responses = iter([
            _Response({"result": {"serverInfo": {"version": "7.1.6"}}}),
            _Response({}),
            _Response({"result": {"tools": [{"name": "search_messages"}, {"name": "aggregate_messages"}]}}),
        ])
        with patch("fgai.graylog_mcp.request.urlopen", side_effect=responses):
            status = GraylogMcpClient("http://graylog/api/mcp", "raw-token").probe()
        self.assertEqual(status["status"], "connected")
        self.assertIn("search_messages", status["tools"])


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