agent_deply/tests/test_mcp_client.py

from pam_deploy_graph.mcp_client import (
    FunctionMcpToolClient,
    HttpMcpToolClient,
    load_mcp_client_config,
    normalize_mcp_tool_list,
    OAuthTokenProvider,
    SessionMcpToolClient,
    StdioMcpToolClient,
    normalize_mcp_sdk_result,
)
from pam_deploy_graph.mcp_factory import build_mcp_runner_from_config
from pam_deploy_graph.mcp_runner import McpActionRunner


def test_function_mcp_client_wraps_callable():
    client = FunctionMcpToolClient(lambda name, args: {"tool": name, "args": args})
    assert client.call_tool("pam_get_online_ips", {"airportCode": "HET"})["tool"] == "pam_get_online_ips"


def test_normalize_mcp_sdk_result_structured_content():
    result = type("Result", (), {"structuredContent": {"ok": True}})()
    assert normalize_mcp_sdk_result(result) == {"ok": True}


def test_session_mcp_client_normalizes_text_json_content():
    content = [type("Text", (), {"text": '{"ok": true}'})()]
    result = type("Result", (), {"content": content})()

    class Session:
        def call_tool(self, tool_name, arguments):
            return result

    client = SessionMcpToolClient(Session())
    assert client.call_tool("tool", {}) == {"ok": True}


def test_normalize_mcp_tool_list():
    result = type(
        "Tools",
        (),
        {"tools": [type("Tool", (), {"name": "pam_get_online_ips"})(), {"name": "verify-ip"}]},
    )()

    assert normalize_mcp_tool_list(result) == ["pam_get_online_ips", "verify-ip"]


def test_load_mcp_client_config(tmp_path):
    path = tmp_path / "mcp.json"
    path.write_text(
        (
            '{"server_name": "pam-node-prod", "transport": "stdio", '
            '"command": "python", "args": ["-m", "server"], '
            '"env": {"PAM_ENV": "test"}, "cwd": "/tmp", "timeout_seconds": 3, '
            '"tool_names": {"get-online-ips": "custom_ips"}}'
        ),
        encoding="utf-8",
    )

    config = load_mcp_client_config(path)

    assert config.server_name == "pam-node-prod"
    assert config.transport == "stdio"
    assert config.command == "python"
    assert config.args == ["-m", "server"]
    assert config.env == {"PAM_ENV": "test"}
    assert config.cwd == "/tmp"
    assert config.timeout_seconds == 3
    assert config.tool_names["get-online-ips"] == "custom_ips"


def test_load_http_mcp_client_config_with_auth(tmp_path):
    path = tmp_path / "mcp.json"
    path.write_text(
        """
        {
          "server_name": "pam-node-prod",
          "transport": "streamable_http",
          "server_url": "https://pam-node.example.com/mcp",
          "auth": {
            "token_url": "https://pam-node-auth.example.com/oauth/token",
            "client_id": "mcp-client",
            "client_secret": "mcp-secret"
          }
        }
        """,
        encoding="utf-8",
    )

    config = load_mcp_client_config(path)

    assert config.transport == "streamable_http"
    assert config.server_url == "https://pam-node.example.com/mcp"
    assert config.auth is not None
    assert config.auth.client_id == "mcp-client"
    assert config.auth.client_secret == "mcp-secret"


def test_build_mcp_runner_from_stdio_config(tmp_path):
    path = tmp_path / "mcp.json"
    path.write_text(
        '{"transport": "stdio", "command": "python", "tool_names": {"verify-ip": "custom_verify"}}',
        encoding="utf-8",
    )

    runner = build_mcp_runner_from_config(path)

    assert isinstance(runner.client, StdioMcpToolClient)
    assert runner.tool_names["verify-ip"] == "custom_verify"


def test_build_mcp_runner_from_http_config(tmp_path):
    path = tmp_path / "mcp.json"
    path.write_text(
        """
        {
          "transport": "sse",
          "server_url": "https://pam-node.example.com/sse",
          "auth": {
            "token_url": "https://pam-node-auth.example.com/oauth/token",
            "client_id": "mcp-client",
            "client_secret": "mcp-secret"
          }
        }
        """,
        encoding="utf-8",
    )

    runner = build_mcp_runner_from_config(path)

    assert isinstance(runner.client, HttpMcpToolClient)
    assert runner.client.transport == "sse"


def test_oauth_token_provider_uses_home_style_form(monkeypatch, tmp_path):
    config = load_mcp_client_config(
        _write_json_config(
            tmp_path,
            {
                "transport": "streamable_http",
                "server_url": "https://pam-node.example.com/mcp",
                "auth": {
                    "token_url": "https://pam-node-auth.example.com/oauth/token",
                    "client_id": "mcp-client",
                    "client_secret": "mcp-secret",
                },
            },
        )
    )
    assert config.auth is not None
    calls = []

    class Response:
        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb):
            return False

        def read(self):
            return b'{"access_token": "token-1", "expires_in": 3600}'

    def fake_urlopen(request, timeout):
        calls.append((request, timeout))
        return Response()

    monkeypatch.setattr("urllib.request.urlopen", fake_urlopen)
    provider = OAuthTokenProvider(config.auth)

    headers = provider.authorization_headers()

    assert headers == {"Authorization": "Bearer token-1"}
    body = calls[0][0].data.decode("utf-8")
    assert "grant_type=client_credentials" in body
    assert "client_id=mcp-client" in body
    assert "client_secret=mcp-secret" in body


def test_mcp_runner_auto_discovers_tool_name():
    class Client:
        def list_tools(self):
            return ["pam_get_online_ips"]

        def call_tool(self, tool_name, arguments):
            return {"IP": ["192.168.1.10"], "COUNT": 1, "TOOL": tool_name}

    runner = McpActionRunner(client=Client())

    result = runner.run("get-online-ips", params={})

    assert result.ok is True
    assert result.tool_name == "pam_get_online_ips"


def _write_json_config(tmpdir, payload):
    path = tmpdir / "mcp.json"
    path.write_text(__import__("json").dumps(payload), encoding="utf-8")
    return str(path)