From bccf6822fb7b083bcdee99c66bfe11e491e5b579 Mon Sep 17 00:00:00 2001 From: Igor Polynets Date: Fri, 9 Jan 2026 20:03:16 -0500 Subject: [PATCH] Refactor MCP server for improved testability Separate implementation functions from MCP tool decorators to enable direct unit testing without MCP framework overhead. The tool wrappers now delegate to _impl functions that contain the actual logic. Co-Authored-By: Claude Opus 4.5 --- telecode/mcp_server.py | 47 +++++++++++++++++++++++++++++++--------- tests/test_mcp_server.py | 44 ++++++++++++++++++------------------- 2 files changed, 59 insertions(+), 32 deletions(-) diff --git a/telecode/mcp_server.py b/telecode/mcp_server.py index b95106a..60895ee 100644 --- a/telecode/mcp_server.py +++ b/telecode/mcp_server.py @@ -86,12 +86,8 @@ def _truncate_message(text: str, limit: int = 3500) -> str: return f"{text[:limit]}\n...[truncated]" -# Initialize FastMCP server -mcp = FastMCP("Telecode MCP Server") - - -@mcp.tool() -def local_claude_code( +# Implementation functions (for testing and reuse) +def _local_claude_code_impl( prompt: str, session_id: Optional[str] = None, timeout_s: Optional[int] = None, @@ -138,8 +134,7 @@ def local_claude_code( return f"Unexpected error: {type(exc).__name__}: {exc}" -@mcp.tool() -def local_codex( +def _local_codex_impl( prompt: str, session_id: Optional[str] = None, timeout_s: Optional[int] = None, @@ -199,8 +194,7 @@ def local_codex( } -@mcp.tool() -def local_cli(command: str, timeout_s: int = 30) -> str: +def _local_cli_impl(command: str, timeout_s: int = 30) -> str: """Execute a shell command locally. Security Warning: @@ -240,6 +234,39 @@ def local_cli(command: str, timeout_s: int = 30) -> str: return f"Error: Command failed: {exc}" +# Initialize FastMCP server +mcp = FastMCP("Telecode MCP Server") + + +# MCP tool wrappers (delegate to implementation functions) +@mcp.tool() +def local_claude_code( + prompt: str, + session_id: Optional[str] = None, + timeout_s: Optional[int] = None, + image_paths: Optional[list[str]] = None +) -> str: + """Execute a prompt using Claude Code CLI.""" + return _local_claude_code_impl(prompt, session_id, timeout_s, image_paths) + + +@mcp.tool() +def local_codex( + prompt: str, + session_id: Optional[str] = None, + timeout_s: Optional[int] = None, + image_paths: Optional[list[str]] = None +) -> dict: + """Execute a prompt using Codex CLI.""" + return _local_codex_impl(prompt, session_id, timeout_s, image_paths) + + +@mcp.tool() +def local_cli(command: str, timeout_s: int = 30) -> str: + """Execute a shell command locally.""" + return _local_cli_impl(command, timeout_s) + + def create_mcp_app(): """Create and configure MCP server app. diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 1975da2..6fca794 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -79,8 +79,8 @@ def mock_run(cmd, **kwargs): monkeypatch.setattr("subprocess.run", mock_run) - from telecode.mcp_server import local_cli - result = local_cli("echo test") + from telecode.mcp_server import _local_cli_impl + result = _local_cli_impl("echo test") assert "test output" in result @@ -91,8 +91,8 @@ def mock_run(cmd, **kwargs): monkeypatch.setattr("subprocess.run", mock_run) - from telecode.mcp_server import local_cli - result = local_cli("sleep 100", timeout_s=10) + from telecode.mcp_server import _local_cli_impl + result = _local_cli_impl("sleep 100", timeout_s=10) assert "timed out" in result.lower() @@ -103,8 +103,8 @@ def mock_run(cmd, **kwargs): monkeypatch.setattr("subprocess.run", mock_run) - from telecode.mcp_server import local_cli - result = local_cli("invalid_command") + from telecode.mcp_server import _local_cli_impl + result = _local_cli_impl("invalid_command") assert "error" in result.lower() or "failed" in result.lower() @@ -115,10 +115,10 @@ def mock_ask_claude_code(prompt, session_id, timeout_s, image_paths): monkeypatch.setattr("telecode.mcp_server.ask_claude_code", mock_ask_claude_code) - from telecode.mcp_server import local_claude_code, _MCP_SESSIONS + from telecode.mcp_server import _local_claude_code_impl, _MCP_SESSIONS _MCP_SESSIONS.clear() - result = local_claude_code("test prompt") + result = _local_claude_code_impl("test prompt") assert result == "Mocked Claude response" @@ -130,8 +130,8 @@ def mock_ask_claude_code(prompt, session_id, timeout_s, image_paths): monkeypatch.setattr("telecode.mcp_server.ask_claude_code", mock_ask_claude_code) - from telecode.mcp_server import local_claude_code - result = local_claude_code("test prompt", session_id="my-session-123") + from telecode.mcp_server import _local_claude_code_impl + result = _local_claude_code_impl("test prompt", session_id="my-session-123") assert result == "Mocked Claude response" @@ -142,8 +142,8 @@ def mock_ask_claude_code(*args, **kwargs): monkeypatch.setattr("telecode.mcp_server.ask_claude_code", mock_ask_claude_code) - from telecode.mcp_server import local_claude_code - result = local_claude_code("test", timeout_s=10) + from telecode.mcp_server import _local_claude_code_impl + result = _local_claude_code_impl("test", timeout_s=10) assert "timed out" in result.lower() @@ -154,8 +154,8 @@ def mock_ask_claude_code(*args, **kwargs): monkeypatch.setattr("telecode.mcp_server.ask_claude_code", mock_ask_claude_code) - from telecode.mcp_server import local_claude_code - result = local_claude_code("test") + from telecode.mcp_server import _local_claude_code_impl + result = _local_claude_code_impl("test") assert "error" in result.lower() assert "session not found" in result.lower() @@ -167,10 +167,10 @@ def mock_ask_codex_exec(prompt, session_id, timeout_s, image_paths): monkeypatch.setattr("telecode.mcp_server.ask_codex_exec", mock_ask_codex_exec) - from telecode.mcp_server import local_codex, _MCP_SESSIONS + from telecode.mcp_server import _local_codex_impl, _MCP_SESSIONS _MCP_SESSIONS.clear() - result = local_codex("test prompt") + result = _local_codex_impl("test prompt") assert isinstance(result, dict) assert result["answer"] == "Mocked answer" assert result["session_id"] == "session-123" @@ -185,8 +185,8 @@ def mock_ask_codex_exec(prompt, session_id, timeout_s, image_paths): monkeypatch.setattr("telecode.mcp_server.ask_codex_exec", mock_ask_codex_exec) - from telecode.mcp_server import local_codex - result = local_codex("test prompt", session_id="my-session-456") + from telecode.mcp_server import _local_codex_impl + result = _local_codex_impl("test prompt", session_id="my-session-456") assert result["session_id"] == "my-session-456" @@ -197,8 +197,8 @@ def mock_ask_codex_exec(*args, **kwargs): monkeypatch.setattr("telecode.mcp_server.ask_codex_exec", mock_ask_codex_exec) - from telecode.mcp_server import local_codex - result = local_codex("test", timeout_s=10) + from telecode.mcp_server import _local_codex_impl + result = _local_codex_impl("test", timeout_s=10) assert isinstance(result, dict) assert "timed out" in result["answer"].lower() @@ -210,8 +210,8 @@ def mock_ask_codex_exec(*args, **kwargs): monkeypatch.setattr("telecode.mcp_server.ask_codex_exec", mock_ask_codex_exec) - from telecode.mcp_server import local_codex - result = local_codex("test") + from telecode.mcp_server import _local_codex_impl + result = _local_codex_impl("test") assert isinstance(result, dict) assert "error" in result["answer"].lower() assert "codex failed" in result["answer"].lower()