improve code coverage

Author: kantord

tests/test_mcp_client.py

@@ -509,6 +509,244 @@ async def test_get_tool_details_connection_error(self, mocker):
         assert "Connection failed" in result["error"]
 
 
+@pytest.mark.asyncio
+class TestCallTool:
+    """Test the call_tool function"""
+
+    async def test_call_tool_success_streamable_http(self, mocker):
+        """Test successful tool call via streamable-http"""
+        workload = {
+            "name": "test-server",
+            "status": "running",
+            "transport_type": "streamable-http",
+            "url": "http://localhost:8080/mcp",
+        }
+
+        mocker.patch("mcp_client.get_workloads", return_value=[workload])
+        mocker.patch(
+            "toolhive_client.discover_toolhive", return_value=("localhost", 8080)
+        )
+
+        # Mock the MCP client
+        mock_result = MagicMock()
+        mock_result.content = [MagicMock(text="tool result")]
+
+        mock_session = MagicMock()
+        mock_session.initialize = AsyncMock()
+        mock_session.call_tool = AsyncMock(return_value=mock_result)
+
+        mock_client_session = MagicMock()
+        mock_client_session.__aenter__ = AsyncMock(return_value=mock_session)
+        mock_client_session.__aexit__ = AsyncMock()
+
+        mock_http = MagicMock()
+        mock_http.__aenter__ = AsyncMock(return_value=("read", "write", lambda: None))
+        mock_http.__aexit__ = AsyncMock()
+
+        mocker.patch("mcp_client.streamablehttp_client", return_value=mock_http)
+        mocker.patch("mcp_client.ClientSession", return_value=mock_client_session)
+
+        result = await mcp_client.call_tool(
+            "test-server", "test_tool", {"param": "value"}
+        )
+
+        assert result == mock_result
+        mock_session.call_tool.assert_called_once_with(
+            "test_tool", arguments={"param": "value"}
+        )
+
+    async def test_call_tool_success_sse(self, mocker):
+        """Test successful tool call via SSE"""
+        workload = {
+            "name": "test-server",
+            "status": "running",
+            "proxy_mode": "sse",
+            "url": "http://localhost:8080/sse",
+        }
+
+        mocker.patch("mcp_client.get_workloads", return_value=[workload])
+        mocker.patch(
+            "toolhive_client.discover_toolhive", return_value=("localhost", 8080)
+        )
+
+        mock_result = MagicMock()
+
+        mock_session = MagicMock()
+        mock_session.initialize = AsyncMock()
+        mock_session.call_tool = AsyncMock(return_value=mock_result)
+
+        mock_client_session = MagicMock()
+        mock_client_session.__aenter__ = AsyncMock(return_value=mock_session)
+        mock_client_session.__aexit__ = AsyncMock()
+
+        mock_sse = MagicMock()
+        mock_sse.__aenter__ = AsyncMock(return_value=("read", "write"))
+        mock_sse.__aexit__ = AsyncMock()
+
+        mocker.patch("mcp_client.sse_client", return_value=mock_sse)
+        mocker.patch("mcp_client.ClientSession", return_value=mock_client_session)
+
+        result = await mcp_client.call_tool(
+            "test-server", "test_tool", {"param": "value"}
+        )
+
+        assert result == mock_result
+
+    async def test_call_tool_workload_not_found(self, mocker):
+        """Test call_tool when workload doesn't exist"""
+        mocker.patch("mcp_client.get_workloads", return_value=[])
+        mocker.patch(
+            "toolhive_client.discover_toolhive", return_value=("localhost", 8080)
+        )
+
+        with pytest.raises(ValueError, match="not found"):
+            await mcp_client.call_tool("nonexistent", "test_tool", {})
+
+    async def test_call_tool_workload_not_running(self, mocker):
+        """Test call_tool when workload is not running"""
+        workload = {
+            "name": "test-server",
+            "status": "stopped",
+            "transport_type": "streamable-http",
+            "url": "http://localhost:8080/mcp",
+        }
+
+        mocker.patch("mcp_client.get_workloads", return_value=[workload])
+        mocker.patch(
+            "toolhive_client.discover_toolhive", return_value=("localhost", 8080)
+        )
+
+        with pytest.raises(RuntimeError, match="not running"):
+            await mcp_client.call_tool("test-server", "test_tool", {})
+
+    async def test_call_tool_no_url(self, mocker):
+        """Test call_tool when workload has no URL"""
+        workload = {
+            "name": "test-server",
+            "status": "running",
+            "transport_type": "streamable-http",
+            "url": "",
+        }
+
+        mocker.patch("mcp_client.get_workloads", return_value=[workload])
+        mocker.patch(
+            "toolhive_client.discover_toolhive", return_value=("localhost", 8080)
+        )
+
+        with pytest.raises(ValueError, match="No URL"):
+            await mcp_client.call_tool("test-server", "test_tool", {})
+
+    async def test_call_tool_unsupported_transport(self, mocker):
+        """Test call_tool with unsupported transport"""
+        workload = {
+            "name": "test-server",
+            "status": "running",
+            "transport_type": "stdio",
+            "url": "http://localhost:8080/mcp",
+        }
+
+        mocker.patch("mcp_client.get_workloads", return_value=[workload])
+        mocker.patch(
+            "toolhive_client.discover_toolhive", return_value=("localhost", 8080)
+        )
+
+        with pytest.raises(ValueError, match="not supported"):
+            await mcp_client.call_tool("test-server", "test_tool", {})
+
+    async def test_call_tool_discovery_fallback(self, mocker):
+        """Test call_tool falls back to defaults when discovery fails"""
+        workload = {
+            "name": "test-server",
+            "status": "running",
+            "transport_type": "streamable-http",
+            "url": "http://localhost:8080/mcp",
+        }
+
+        mock_get_workloads = mocker.patch(
+            "mcp_client.get_workloads", return_value=[workload]
+        )
+        mocker.patch(
+            "toolhive_client.discover_toolhive",
+            side_effect=Exception("Discovery failed"),
+        )
+
+        mock_result = MagicMock()
+        mock_session = MagicMock()
+        mock_session.initialize = AsyncMock()
+        mock_session.call_tool = AsyncMock(return_value=mock_result)
+
+        mock_client_session = MagicMock()
+        mock_client_session.__aenter__ = AsyncMock(return_value=mock_session)
+        mock_client_session.__aexit__ = AsyncMock()
+
+        mock_http = MagicMock()
+        mock_http.__aenter__ = AsyncMock(return_value=("read", "write", lambda: None))
+        mock_http.__aexit__ = AsyncMock()
+
+        mocker.patch("mcp_client.streamablehttp_client", return_value=mock_http)
+        mocker.patch("mcp_client.ClientSession", return_value=mock_client_session)
+
+        # Should not raise, should fall back to defaults
+        result = await mcp_client.call_tool("test-server", "test_tool", {})
+
+        assert result == mock_result
+        # Should have been called with default host/port
+        mock_get_workloads.assert_called_once_with("127.0.0.1", 8080)
+
+
+@pytest.mark.asyncio
+class TestGetWorkloadsUrlRewriting:
+    """Test localhost URL rewriting for container networking"""
+
+    async def test_rewrites_localhost_urls(self, mocker):
+        """Test that localhost URLs are rewritten to use the actual host"""
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "workloads": [
+                {
+                    "name": "workload1",
+                    "url": "http://localhost:9000/mcp",
+                },
+                {
+                    "name": "workload2",
+                    "url": "http://127.0.0.1:9001/sse",
+                },
+            ]
+        }
+
+        mock_client = MagicMock()
+        mock_client.__aenter__.return_value.get = AsyncMock(return_value=mock_response)
+        mocker.patch("httpx.AsyncClient", return_value=mock_client)
+
+        # Call with a different host (simulating container environment)
+        result = await mcp_client.get_workloads(host="192.168.1.100", port=8080)
+
+        # URLs should be rewritten
+        assert result[0]["url"] == "http://192.168.1.100:9000/mcp"
+        assert result[1]["url"] == "http://192.168.1.100:9001/sse"
+
+    async def test_preserves_non_localhost_urls(self, mocker):
+        """Test that non-localhost URLs are not rewritten"""
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "workloads": [
+                {
+                    "name": "workload1",
+                    "url": "http://some-service:9000/mcp",
+                },
+            ]
+        }
+
+        mock_client = MagicMock()
+        mock_client.__aenter__.return_value.get = AsyncMock(return_value=mock_response)
+        mocker.patch("httpx.AsyncClient", return_value=mock_client)
+
+        result = await mcp_client.get_workloads(host="192.168.1.100", port=8080)
+
+        # URL should not be rewritten
+        assert result[0]["url"] == "http://some-service:9000/mcp"
+
+
 @pytest.mark.asyncio
 class TestSelfFiltering:
     """Test that mcp-shell filters itself out from tool listings"""

tests/test_shell_engine.py

@@ -145,6 +145,21 @@ async def test_shell_stage_for_each_mode(self):
         assert "apricot" in output
         assert "banana" not in output
 
+    async def test_shell_stage_for_each_without_trailing_newline(self):
+        """Test shell command with for_each when input lacks trailing newline."""
+        mock_caller = AsyncMock()
+        engine = ShellEngine(tool_caller=mock_caller)
+
+        # Input without trailing newline - last line should still be processed
+        upstream = iter(["apple\nbanana\napricot"])
+        result = list(engine.shell_stage("grep", ["^a"], upstream, for_each=True))
+
+        output = "".join(result)
+        # All lines starting with 'a' should be processed, including 'apricot'
+        assert "apple" in output
+        assert "apricot" in output
+        assert "banana" not in output
+
     async def test_shell_stage_empty_input(self):
         """Test shell command with empty input."""
         mock_caller = AsyncMock()
@@ -263,6 +278,72 @@ async def test_tool_stage_result_with_no_content_attribute(self):
 
         assert result == "plain string result"
 
+    async def test_tool_stage_for_each_without_trailing_newline(self):
+        """Test tool call with for_each when input doesn't have trailing newline."""
+        mock_caller = AsyncMock(return_value=MockToolResult("result"))
+        engine = ShellEngine(tool_caller=mock_caller)
+
+        # Input without trailing newline - should still process the last line
+        jsonl_input = '{"url": "http://example.com/1"}\n{"url": "http://example.com/2"}'
+        upstream = iter([jsonl_input])
+
+        result = await engine.tool_stage(
+            "test_server", "fetch", {}, upstream, for_each=True
+        )
+
+        # Should be called twice (once per line, including line without newline)
+        assert mock_caller.call_count == 2
+        assert "result" in result
+
+    async def test_tool_stage_non_dict_json_upstream(self):
+        """Test tool_stage with non-dict JSON upstream (array) in non-for_each mode."""
+        mock_caller = AsyncMock(return_value=MockToolResult("result"))
+        engine = ShellEngine(tool_caller=mock_caller)
+
+        # JSON array as upstream
+        array_input = '["item1", "item2", "item3"]'
+        upstream = iter([array_input])
+
+        result = await engine.tool_stage("test_server", "test_tool", {}, upstream)
+
+        # Should add array as 'input' field
+        mock_caller.assert_called_once()
+        call_args = mock_caller.call_args[0][2]
+        assert call_args["input"] == ["item1", "item2", "item3"]
+
+    async def test_tool_stage_plain_text_upstream(self):
+        """Test tool_stage with plain text (non-JSON) upstream."""
+        mock_caller = AsyncMock(return_value=MockToolResult("result"))
+        engine = ShellEngine(tool_caller=mock_caller)
+
+        # Plain text that isn't valid JSON
+        text_input = "some plain text data"
+        upstream = iter([text_input])
+
+        result = await engine.tool_stage("test_server", "test_tool", {}, upstream)
+
+        # Should add text as 'input' field
+        mock_caller.assert_called_once()
+        call_args = mock_caller.call_args[0][2]
+        assert call_args["input"] == "some plain text data"
+
+    async def test_tool_stage_non_dict_json_does_not_override_existing_input(self):
+        """Test that non-dict JSON doesn't override explicit 'input' arg."""
+        mock_caller = AsyncMock(return_value=MockToolResult("result"))
+        engine = ShellEngine(tool_caller=mock_caller)
+
+        array_input = '["upstream_data"]'
+        upstream = iter([array_input])
+
+        result = await engine.tool_stage(
+            "test_server", "test_tool", {"input": "explicit_input"}, upstream
+        )
+
+        # Explicit input should be preserved
+        mock_caller.assert_called_once()
+        call_args = mock_caller.call_args[0][2]
+        assert call_args["input"] == "explicit_input"
+
 
 @pytest.mark.asyncio
 class TestExecutePipeline:
@@ -362,8 +443,8 @@ async def test_execute_pipeline_missing_command_field(self):
         assert "Pipeline execution failed" in result
         assert "missing 'command' field" in result
 
-    async def test_execute_pipeline_missing_tool_fields(self):
-        """Test pipeline with missing tool fields."""
+    async def test_execute_pipeline_missing_tool_name_field(self):
+        """Test pipeline with missing tool name field."""
         mock_caller = AsyncMock()
         engine = ShellEngine(tool_caller=mock_caller)
 
@@ -375,6 +456,19 @@ async def test_execute_pipeline_missing_tool_fields(self):
         assert "Pipeline execution failed" in result
         assert "missing 'name' field" in result
 
+    async def test_execute_pipeline_missing_tool_server_field(self):
+        """Test pipeline with missing tool server field."""
+        mock_caller = AsyncMock()
+        engine = ShellEngine(tool_caller=mock_caller)
+
+        # Missing "server" field
+        pipeline = [{"type": "tool", "name": "test_tool", "args": {}}]
+
+        result = await engine.execute_pipeline(pipeline)
+
+        assert "Pipeline execution failed" in result
+        assert "missing 'server' field" in result
+
     async def test_execute_pipeline_invalid_args_type(self):
         """Test pipeline with invalid args type (not a list)."""
         mock_caller = AsyncMock()