refactor(__init__.py): update module imports to only include main and add module docstring for clarity

refactor(server.py): simplify error handling in command execution by raising an error directly instead of appending to content
test: enhance test coverage for tool execution with additional test cases for valid and invalid commands

Author: tumf

mcp_shell_server/__init__.py

@@ -1,5 +1,5 @@
-from .server import ShellServer, main
-from .shell_executor import ShellExecutor
+"""MCP Shell Server Package."""
+from .server import main
 
 __version__ = "0.1.0"
-__all__ = ["ShellServer", "ShellExecutor", "main"]
+__all__ = ["main"]

mcp_shell_server/server.py

@@ -52,14 +52,13 @@ async def run_tool(self, arguments: dict) -> Sequence[TextContent]:
 
         result = await self.executor.execute(command, stdin)
 
+        # Raise error if command execution failed
+        if result.get("error"):
+            raise RuntimeError(result["error"])
+
         # Convert executor result to TextContent sequence
         content: list[TextContent] = []
 
-        if result.get("error"):
-            content.append(TextContent(
-                type="text",
-                text=result["error"]
-            ))
         if result.get("stdout"):
             content.append(TextContent(
                 type="text",

tests/test_server.py

@@ -1,41 +1,82 @@
 import pytest
-from mcp_shell_server.server import ShellServer
+from mcp.types import TextContent, Tool
+from mcp_shell_server.server import call_tool, list_tools, tool_handler
 
-@pytest.fixture
-def server():
-    return ShellServer()
 
 @pytest.mark.asyncio
-async def test_server_handle_empty_request(server):
-    response = await server.handle({})
-    assert response["error"] == "No command provided"
-    assert response["status"] == 1
+async def test_list_tools():
+    """Test listing of available tools"""
+    tools = await list_tools()
+    assert len(tools) == 1
+    tool = tools[0]
+    assert isinstance(tool, Tool)
+    assert tool.name == "execute"
+    assert tool.description
+    assert tool.inputSchema["type"] == "object"
+    assert "command" in tool.inputSchema["properties"]
+    assert "stdin" in tool.inputSchema["properties"]
+    assert tool.inputSchema["required"] == ["command"]
+
 
 @pytest.mark.asyncio
-async def test_server_handle_valid_command(server, monkeypatch):
+async def test_call_tool_valid_command(monkeypatch):
+    """Test execution of a valid command"""
     monkeypatch.setenv("ALLOW_COMMANDS", "echo")
-    response = await server.handle({
+    result = await call_tool("execute", {
         "command": ["echo", "hello world"]
     })
-    assert response["stdout"].strip() == "hello world"
-    assert response["status"] == 0
-    assert "execution_time" in response
+    assert len(result) == 1
+    assert isinstance(result[0], TextContent)
+    assert result[0].type == "text"
+    assert result[0].text.strip() == "hello world"
+
 
 @pytest.mark.asyncio
-async def test_server_handle_stdin(server, monkeypatch):
+async def test_call_tool_with_stdin(monkeypatch):
+    """Test command execution with stdin"""
     monkeypatch.setenv("ALLOW_COMMANDS", "cat")
-    response = await server.handle({
+    result = await call_tool("execute", {
         "command": ["cat"],
         "stdin": "test input"
     })
-    assert response["stdout"].strip() == "test input"
-    assert response["status"] == 0
+    assert len(result) == 1
+    assert isinstance(result[0], TextContent)
+    assert result[0].type == "text"
+    assert result[0].text.strip() == "test input"
+
 
 @pytest.mark.asyncio
-async def test_server_handle_invalid_command(server, monkeypatch):
+async def test_call_tool_invalid_command(monkeypatch):
+    """Test execution of an invalid command"""
     monkeypatch.setenv("ALLOW_COMMANDS", "echo")
-    response = await server.handle({
-        "command": ["invalid_command"]
-    })
-    assert response["error"] == "Command not allowed: invalid_command"
-    assert response["status"] == 1
+    with pytest.raises(RuntimeError) as excinfo:
+        await call_tool("execute", {
+            "command": ["invalid_command"]
+        })
+    assert "Command not allowed: invalid_command" in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_call_tool_unknown_tool():
+    """Test calling an unknown tool"""
+    with pytest.raises(RuntimeError) as excinfo:
+        await call_tool("unknown_tool", {})
+    assert "Unknown tool: unknown_tool" in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_call_tool_invalid_arguments():
+    """Test calling a tool with invalid arguments"""
+    with pytest.raises(RuntimeError) as excinfo:
+        await call_tool("execute", "not a dict")
+    assert "Arguments must be a dictionary" in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_call_tool_empty_command():
+    """Test execution with empty command"""
+    with pytest.raises(RuntimeError) as excinfo:
+        await call_tool("execute", {
+            "command": []
+        })
+    assert "No command provided" in str(excinfo.value)