400 lines
14 KiB
Python
400 lines
14 KiB
Python
"""Tests for MCP client and McpStep execution."""
|
|
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from cmdforge.mcp_client import (
|
|
DEFAULT_INHERITED_ENV,
|
|
McpServerConfig,
|
|
McpClientManager,
|
|
_build_server_env,
|
|
_normalize_result,
|
|
_sanitize,
|
|
_fingerprint,
|
|
load_mcp_config,
|
|
save_mcp_config,
|
|
)
|
|
from cmdforge.runner import _substitute_mcp_args, _deep_substitute, run_tool
|
|
from cmdforge.tool import McpStep, Tool
|
|
|
|
|
|
class TestMcpServerConfig:
|
|
def test_defaults(self):
|
|
cfg = McpServerConfig(name="test")
|
|
assert cfg.transport == "stdio"
|
|
assert cfg.timeout == 30
|
|
assert cfg.args == []
|
|
assert cfg.env == {}
|
|
assert cfg.approved is False
|
|
|
|
@pytest.mark.parametrize(
|
|
"change, message",
|
|
[
|
|
({"transport": "streamable-http"}, "Unsupported MCP transport"),
|
|
({"args": "-y"}, "list of strings"),
|
|
({"timeout": 0}, "greater than 0"),
|
|
({"timeout": True}, "must be a number"),
|
|
],
|
|
)
|
|
def test_validation_rejects_invalid_config(self, change, message):
|
|
cfg = McpServerConfig(name="test", command="server", **change)
|
|
with pytest.raises(ValueError, match=message):
|
|
cfg.validate()
|
|
|
|
def test_fingerprint_changes_with_command(self):
|
|
a = McpServerConfig(name="a", command="cmd-a")
|
|
b = McpServerConfig(name="b", command="cmd-b")
|
|
assert _fingerprint(a) != _fingerprint(b)
|
|
|
|
def test_fingerprint_same_for_same_config(self):
|
|
a = McpServerConfig(name="a", command="cmd", args=["-v"])
|
|
b = McpServerConfig(name="b", command="cmd", args=["-v"])
|
|
assert _fingerprint(a) == _fingerprint(b)
|
|
|
|
|
|
class TestSanitize:
|
|
def test_redacts_bearer(self):
|
|
assert "redacted" in _sanitize("Bearer abcdefghijklmnopqrstuvwxyz")
|
|
|
|
def test_redacts_sk_prefix(self):
|
|
assert "redacted" in _sanitize("sk-abcdefghijklmnopqrstuvwxyz")
|
|
|
|
def test_preserves_normal_text(self):
|
|
assert _sanitize("hello world") == "hello world"
|
|
|
|
def test_redacts_explicit_secret(self):
|
|
assert _sanitize("failed with hunter2", ["hunter2"]) == "failed with [redacted]"
|
|
|
|
|
|
class TestEnvironmentIsolation:
|
|
def test_only_allowlisted_environment_is_inherited(self, monkeypatch):
|
|
monkeypatch.setenv("CMDFORGE_TEST_SECRET", "do-not-leak")
|
|
cfg = McpServerConfig(name="test", command="server")
|
|
environment = _build_server_env(cfg)
|
|
assert "CMDFORGE_TEST_SECRET" not in environment
|
|
assert set(environment).issubset(set(DEFAULT_INHERITED_ENV))
|
|
|
|
def test_explicit_environment_reference_is_resolved(self, monkeypatch):
|
|
monkeypatch.setenv("CMDFORGE_TEST_SECRET", "allowed")
|
|
cfg = McpServerConfig(
|
|
name="test", command="server", env={"SERVER_TOKEN": "${CMDFORGE_TEST_SECRET}"}
|
|
)
|
|
assert _build_server_env(cfg)["SERVER_TOKEN"] == "allowed"
|
|
|
|
def test_missing_environment_reference_fails(self, monkeypatch):
|
|
monkeypatch.delenv("CMDFORGE_MISSING", raising=False)
|
|
cfg = McpServerConfig(
|
|
name="test", command="server", env={"SERVER_TOKEN": "${CMDFORGE_MISSING}"}
|
|
)
|
|
with pytest.raises(ValueError, match="is not set"):
|
|
_build_server_env(cfg)
|
|
|
|
|
|
class TestMcpConfigPersistence:
|
|
@pytest.fixture
|
|
def temp_mcp_file(self, tmp_path, monkeypatch):
|
|
config_file = tmp_path / "mcp.yaml"
|
|
monkeypatch.setattr("cmdforge.mcp_client.MCP_CONFIG_FILE", config_file)
|
|
yield config_file
|
|
|
|
def test_save_and_load(self, temp_mcp_file):
|
|
servers = [
|
|
McpServerConfig(
|
|
name="filesystem",
|
|
command="npx",
|
|
args=["-y", "@scope/server", "/tmp"],
|
|
timeout=30,
|
|
approved=True,
|
|
)
|
|
]
|
|
save_mcp_config(servers)
|
|
loaded = load_mcp_config()
|
|
assert len(loaded) == 1
|
|
assert loaded[0].name == "filesystem"
|
|
assert loaded[0].command == "npx"
|
|
assert loaded[0].args == ["-y", "@scope/server", "/tmp"]
|
|
assert loaded[0].approved is True
|
|
|
|
def test_load_no_file(self, temp_mcp_file):
|
|
assert load_mcp_config() == []
|
|
|
|
def test_save_creates_0600_permissions(self, temp_mcp_file):
|
|
save_mcp_config([McpServerConfig(name="test", command="echo")])
|
|
perms = oct(temp_mcp_file.stat().st_mode & 0o777)
|
|
assert perms == "0o600"
|
|
|
|
|
|
class TestResultNormalization:
|
|
def test_auto_prefers_structured_content(self):
|
|
class FakeResult:
|
|
isError = False
|
|
structuredContent = {"sum": 3}
|
|
content = [type("Block", (), {"text": "plain text"})()]
|
|
|
|
result = _normalize_result(FakeResult(), "auto")
|
|
assert result == {"sum": 3}
|
|
|
|
def test_auto_falls_back_to_text(self):
|
|
class FakeResult:
|
|
isError = False
|
|
structuredContent = None
|
|
content = [type("Block", (), {"type": "text", "text": "hello"})()]
|
|
|
|
result = _normalize_result(FakeResult(), "auto")
|
|
assert result == "hello"
|
|
|
|
def test_auto_preserves_mixed_content(self):
|
|
class FakeResult:
|
|
isError = False
|
|
structuredContent = None
|
|
content = [
|
|
{"type": "text", "text": "hello"},
|
|
{"type": "image", "mimeType": "image/png", "data": "AA=="},
|
|
]
|
|
|
|
result = _normalize_result(FakeResult(), "auto")
|
|
assert result == FakeResult.content
|
|
|
|
def test_content_mode_always_returns_blocks(self):
|
|
class FakeResult:
|
|
isError = False
|
|
structuredContent = None
|
|
content = [{"type": "text", "text": "hello"}]
|
|
|
|
assert _normalize_result(FakeResult(), "content") == FakeResult.content
|
|
|
|
def test_structured_fails_without_structured_content(self):
|
|
class FakeResult:
|
|
isError = False
|
|
structuredContent = None
|
|
content = [type("Block", (), {"type": "text", "text": "hello"})()]
|
|
|
|
with pytest.raises(ValueError, match="structuredContent"):
|
|
_normalize_result(FakeResult(), "structured")
|
|
|
|
def test_text_mode(self):
|
|
class FakeResult:
|
|
isError = False
|
|
structuredContent = None
|
|
content = [
|
|
type("Block", (), {"type": "text", "text": "line1"})(),
|
|
type("Block", (), {"type": "text", "text": "line2"})(),
|
|
]
|
|
|
|
result = _normalize_result(FakeResult(), "text")
|
|
assert result == "line1\nline2"
|
|
|
|
def test_iserror_causes_failure(self):
|
|
class FakeResult:
|
|
isError = True
|
|
structuredContent = None
|
|
content = [type("Block", (), {"type": "text", "text": "error message"})()]
|
|
|
|
with pytest.raises(RuntimeError, match="error"):
|
|
_normalize_result(FakeResult(), "auto")
|
|
|
|
|
|
class TestMcpClientManager:
|
|
@pytest.fixture
|
|
def manager(self, tmp_path, monkeypatch):
|
|
config_file = tmp_path / "mcp.yaml"
|
|
monkeypatch.setattr("cmdforge.mcp_client.MCP_CONFIG_FILE", config_file)
|
|
return McpClientManager()
|
|
|
|
def test_list_servers_empty(self, manager):
|
|
assert manager.list_servers() == []
|
|
|
|
def test_list_servers_with_config(self, manager, tmp_path, monkeypatch):
|
|
config_file = tmp_path / "mcp.yaml"
|
|
monkeypatch.setattr("cmdforge.mcp_client.MCP_CONFIG_FILE", config_file)
|
|
save_mcp_config([McpServerConfig(name="test", command="echo")])
|
|
|
|
mgr = McpClientManager()
|
|
servers = mgr.list_servers()
|
|
assert len(servers) == 1
|
|
assert servers[0].name == "test"
|
|
|
|
def test_call_tool_unknown_server(self, manager):
|
|
with pytest.raises(KeyError, match="not configured"):
|
|
manager.call_tool("unknown", "tool", {})
|
|
|
|
def test_unapproved_server_is_not_executed(self, tmp_path, monkeypatch):
|
|
config_file = tmp_path / "mcp.yaml"
|
|
monkeypatch.setattr("cmdforge.mcp_client.MCP_CONFIG_FILE", config_file)
|
|
save_mcp_config([McpServerConfig(name="test", command="echo")])
|
|
with pytest.raises(PermissionError, match="not approved"):
|
|
McpClientManager().discover("test")
|
|
|
|
def test_missing_sdk_has_actionable_error(self, tmp_path, monkeypatch):
|
|
config_file = tmp_path / "mcp.yaml"
|
|
monkeypatch.setattr("cmdforge.mcp_client.MCP_CONFIG_FILE", config_file)
|
|
save_mcp_config([
|
|
McpServerConfig(name="test", command="echo", approved=True)
|
|
])
|
|
monkeypatch.setitem(sys.modules, "mcp", None)
|
|
with pytest.raises(ImportError, match=r"cmdforge\[mcp\]"):
|
|
McpClientManager().discover("test")
|
|
|
|
|
|
class TestArgumentSubstitution:
|
|
def test_single_variable_preserves_integer(self):
|
|
result = _deep_substitute("42", {"limit": 42})
|
|
assert result == "42" # no braces, so no substitution
|
|
|
|
def test_braced_variable_preserves_type(self):
|
|
variables = {"settings": {"limit": 50}}
|
|
result = _substitute_mcp_args(
|
|
{"options": {"limit": "{settings.limit}"}}, variables
|
|
)
|
|
assert result["options"]["limit"] == 50
|
|
|
|
def test_nested_dict_substitution(self):
|
|
variables = {"path": "/tmp/data.csv"}
|
|
result = _substitute_mcp_args(
|
|
{"files": [{"source": "{path}"}]}, variables
|
|
)
|
|
assert result["files"][0]["source"] == "/tmp/data.csv"
|
|
|
|
def test_string_with_variable_inside_text(self):
|
|
variables = {"name": "report"}
|
|
result = _deep_substitute("/tmp/{name}.csv", variables)
|
|
assert result == "/tmp/report.csv"
|
|
|
|
@pytest.mark.parametrize(
|
|
"value",
|
|
[
|
|
{"a": 1},
|
|
[1, 2],
|
|
-5,
|
|
"001",
|
|
"true",
|
|
False,
|
|
None,
|
|
],
|
|
)
|
|
def test_exact_reference_preserves_original_type(self, value):
|
|
assert _deep_substitute("{value}", {"value": value}) == value
|
|
|
|
|
|
class TestMcpStep:
|
|
def test_round_trip(self):
|
|
original = McpStep(
|
|
server="fixture", tool="add", arguments={"a": 1, "b": 2},
|
|
output_var="sum", result_mode="structured",
|
|
)
|
|
assert McpStep.from_dict(original.to_dict()) == original
|
|
|
|
@pytest.mark.parametrize(
|
|
"field, value",
|
|
[("server", ""), ("tool", None), ("arguments", []), ("result_mode", "raw")],
|
|
)
|
|
def test_rejects_invalid_fields(self, field, value):
|
|
data = {"server": "fixture", "tool": "echo", field: value}
|
|
with pytest.raises(ValueError):
|
|
McpStep.from_dict(data)
|
|
|
|
def test_runner_preserves_structured_result(self, monkeypatch):
|
|
calls = []
|
|
|
|
class FakeManager:
|
|
def call_tool(self, server, tool, arguments, result_mode="auto"):
|
|
calls.append((server, tool, arguments, result_mode))
|
|
return {"sum": 3}
|
|
|
|
monkeypatch.setattr("cmdforge.runner.McpClientManager", FakeManager)
|
|
tool = Tool(
|
|
name="mcp-add",
|
|
steps=[
|
|
McpStep(
|
|
server="fixture",
|
|
tool="add",
|
|
arguments={"a": "{input}", "b": 2},
|
|
output_var="result",
|
|
)
|
|
],
|
|
output="{result.sum}",
|
|
)
|
|
output, exit_code = run_tool(tool, 1, {})
|
|
assert (output, exit_code) == ("3", 0)
|
|
assert calls == [("fixture", "add", {"a": 1, "b": 2}, "auto")]
|
|
|
|
|
|
class TestRealMcpSdk:
|
|
@pytest.fixture
|
|
def manager(self, tmp_path, monkeypatch):
|
|
pytest.importorskip("mcp")
|
|
config_file = tmp_path / "mcp.yaml"
|
|
monkeypatch.setattr("cmdforge.mcp_client.MCP_CONFIG_FILE", config_file)
|
|
fixture = Path(__file__).parent / "fixtures" / "mcp_test_server.py"
|
|
save_mcp_config([
|
|
McpServerConfig(
|
|
name="fixture",
|
|
command=sys.executable,
|
|
args=[str(fixture)],
|
|
timeout=10,
|
|
approved=True,
|
|
)
|
|
])
|
|
return McpClientManager()
|
|
|
|
def test_discovery_uses_mcp_handshake_and_tools_list(self, manager):
|
|
tools = manager.discover("fixture")
|
|
schemas = {tool["name"]: tool for tool in tools}
|
|
assert {"echo", "add", "read_env", "current_directory", "wait"}.issubset(schemas)
|
|
assert schemas["add"]["inputSchema"]["required"] == ["a", "b"]
|
|
|
|
def test_calls_text_and_structured_tools(self, manager):
|
|
assert manager.call_tool("fixture", "echo", {"message": "hello"}, "text") == "hello"
|
|
assert manager.call_tool("fixture", "add", {"a": 1, "b": 2}) == {"sum": 3}
|
|
|
|
def test_parent_secret_is_not_leaked(self, manager, monkeypatch):
|
|
monkeypatch.setenv("CMDFORGE_TEST_SECRET", "do-not-leak")
|
|
result = manager.call_tool(
|
|
"fixture", "read_env", {"name": "CMDFORGE_TEST_SECRET"}, "text"
|
|
)
|
|
assert result == ""
|
|
|
|
def test_configured_working_directory_is_used(self, manager, tmp_path):
|
|
manager.list_servers()
|
|
manager._configs["fixture"].cwd = str(tmp_path)
|
|
assert manager.call_tool("fixture", "current_directory", {}, "text") == str(tmp_path)
|
|
|
|
def test_timeout_is_enforced(self, manager):
|
|
manager.list_servers()
|
|
manager._configs["fixture"].timeout = 0.5
|
|
started = time.monotonic()
|
|
with pytest.raises(TimeoutError, match="timed out"):
|
|
manager.call_tool("fixture", "wait", {"seconds": 10}, "text")
|
|
assert time.monotonic() - started < 4
|
|
|
|
|
|
class TestMcpCli:
|
|
def test_add_accepts_leading_dash_argument(self, tmp_path, monkeypatch, capsys):
|
|
from cmdforge.cli import main
|
|
|
|
config_file = tmp_path / "mcp.yaml"
|
|
monkeypatch.setattr("cmdforge.mcp_client.MCP_CONFIG_FILE", config_file)
|
|
monkeypatch.setattr(
|
|
sys,
|
|
"argv",
|
|
[
|
|
"cmdforge", "mcp", "add", "fixture", "--command", "npx",
|
|
"--arg=-y", "--arg", "@scope/server",
|
|
],
|
|
)
|
|
assert main() == 0
|
|
loaded = load_mcp_config()
|
|
assert loaded[0].args == ["-y", "@scope/server"]
|
|
assert loaded[0].approved is True
|
|
assert "saved and approved" in capsys.readouterr().out
|
|
|
|
def test_bare_mcp_lists_servers(self, tmp_path, monkeypatch, capsys):
|
|
from cmdforge.cli import main
|
|
|
|
monkeypatch.setattr("cmdforge.mcp_client.MCP_CONFIG_FILE", tmp_path / "mcp.yaml")
|
|
monkeypatch.setattr(sys, "argv", ["cmdforge", "mcp"])
|
|
assert main() == 0
|
|
assert "No MCP servers configured" in capsys.readouterr().out
|