MCP Prompt Cleaner

from unittest.mock import patch, mock_open, AsyncMock import pytest from main import list_tools class TestIntegration: """Integration tests for MCP server functionality""" @pytest.fixture def mock_system_prompt(self): """Mock system prompt content""" return """[ROLE] "Prompt Cleaner"; Expert prompt engineer [TASK] Refine RAW_PROMPT; preserve intent, placeholders, files, code callouts [INPUT] RAW_PROMPT: string [OUTPUT] STRICT; stdout = Valid JSON (UTF-8), double quotes, no trailing commas, no extra text; FIELDS: { "cleaned": "string", "notes": ["string"], "open_questions": ["string"], "risks": ["string"], "unchanged": boolean, "quality": { "score": integer (1-5), "reasons": ["string"] } } [QUALITY_GATE] Score RAW_PROMPT 0–5 (1pt each): intent clear; io stated/N/A; constraints/acceptance present/N/A; no contradictions; If score ≥4 AND no redactions: unchanged=true and cleaned=RAW_PROMPT (byte-exact). Else unchanged=false and refine [CLEANING_RULES] Concise, actionable, unambiguous; Use "\n- " for lists; specify inputs/outputs when present or clearly implied; Developer tone if code/spec; include types and edge/error cases; Don't invent requirements or change scope; preserve {{var}}, <VAR>, $VAR, backticks; Keep original language [TROUBLESHOOT] If RAW_PROMPT is a direct question or short-answer: produce normal json preserving original prompt [MODE_CONTEXT] Consider MODE (code/general) and CONTEXT for domain-specific improvements""" @pytest.fixture def mock_llm_response(self): """Mock LLM response""" return '{"cleaned": "Enhanced prompt", "notes": ["Added specificity"], "open_questions": [], "risks": [], "unchanged": false, "quality": {"score": 4, "reasons": ["Clear and actionable"]}}' @pytest.mark.asyncio async def test_list_tools(self): """Test tool listing functionality""" tools = await list_tools() assert len(tools) == 1 assert tools[0].name == "clean_prompt" assert "Enhance and clean raw prompts" in tools[0].description assert hasattr(tools[0], "inputSchema") # Verify input schema structure schema = tools[0].inputSchema assert schema["type"] == "object" assert "properties" in schema assert "raw_prompt" in schema["properties"] assert "context" in schema["properties"] assert "mode" in schema["properties"] assert "temperature" in schema["properties"] @pytest.mark.asyncio async def test_clean_prompt_tool_call(self, mock_system_prompt, mock_llm_response): """Test clean_prompt tool call through MCP protocol""" with patch("builtins.open", mock_open(read_data=mock_system_prompt)): with patch("tools.cleaner.LLMClient") as mock_client_class: mock_client = AsyncMock() mock_client.chat_completions.return_value = mock_llm_response mock_client_class.return_value = mock_client # Test the tool function directly from main import clean_prompt_tool result = await clean_prompt_tool( raw_prompt="help me write code", context="web development", mode="general", temperature=0.2, ) assert result.cleaned == "Enhanced prompt" assert result.notes == ["Added specificity"] assert result.unchanged is False assert result.quality.score == 4 @pytest.mark.asyncio async def test_clean_prompt_defaults(self, mock_system_prompt, mock_llm_response): """Test clean_prompt tool with default parameters""" with patch("builtins.open", mock_open(read_data=mock_system_prompt)): with patch("tools.cleaner.LLMClient") as mock_client_class: mock_client = AsyncMock() mock_client.chat_completions.return_value = mock_llm_response mock_client_class.return_value = mock_client # Test the tool function directly from main import clean_prompt_tool result = await clean_prompt_tool( raw_prompt="fix this bug", context="Python code" ) assert result.cleaned == "Enhanced prompt" assert "cleaned" in result.model_dump() @pytest.mark.asyncio async def test_clean_prompt_minimal_input( self, mock_system_prompt, mock_llm_response ): """Test clean_prompt tool with minimal required input""" with patch("builtins.open", mock_open(read_data=mock_system_prompt)): with patch("tools.cleaner.LLMClient") as mock_client_class: mock_client = AsyncMock() mock_client.chat_completions.return_value = mock_llm_response mock_client_class.return_value = mock_client # Test the tool function directly from main import clean_prompt_tool result = await clean_prompt_tool(raw_prompt="make this better") assert result.cleaned == "Enhanced prompt" assert "cleaned" in result.model_dump() @pytest.mark.asyncio async def test_input_validation_error(self): """Test input validation error handling""" from main import clean_prompt_tool # Test with invalid input - should raise ValueError with pytest.raises(ValueError): await clean_prompt_tool( raw_prompt="", # Invalid: empty string mode="invalid", # Invalid: not code or general ) @pytest.mark.asyncio async def test_llm_error_handling(self, mock_system_prompt): """Test LLM error handling through MCP protocol""" with patch("builtins.open", mock_open(read_data=mock_system_prompt)): with patch("tools.cleaner.LLMClient") as mock_client_class: mock_client = AsyncMock() mock_client.chat_completions.side_effect = Exception("LLM API error") mock_client_class.return_value = mock_client from main import clean_prompt_tool # Should raise the exception with pytest.raises(Exception, match="LLM API error"): await clean_prompt_tool(raw_prompt="test prompt") @pytest.mark.asyncio async def test_minimal_input(self, mock_system_prompt, mock_llm_response): """Test with minimal required input""" with patch("builtins.open", mock_open(read_data=mock_system_prompt)): with patch("tools.cleaner.LLMClient") as mock_client_class: mock_client = AsyncMock() mock_client.chat_completions.return_value = mock_llm_response mock_client_class.return_value = mock_client # Test with only required raw_prompt from main import clean_prompt_tool result = await clean_prompt_tool(raw_prompt="help me") assert result.cleaned == "Enhanced prompt" assert "cleaned" in result.model_dump() @pytest.mark.asyncio async def test_code_mode_integration(self, mock_system_prompt, mock_llm_response): """Test code mode through integration""" with patch("builtins.open", mock_open(read_data=mock_system_prompt)): with patch("tools.cleaner.LLMClient") as mock_client_class: mock_client = AsyncMock() mock_client.chat_completions.return_value = mock_llm_response mock_client_class.return_value = mock_client from main import clean_prompt_tool result = await clean_prompt_tool( raw_prompt="write a function", context="Python development", mode="code", temperature=0.1, ) # Verify the LLM was called with correct parameters mock_client.chat_completions.assert_called_once() call_args = mock_client.chat_completions.call_args messages = call_args[1]["messages"] assert "MODE: code" in messages[1]["content"] assert call_args[1]["temperature"] == 0.1 # Verify response format assert result.cleaned == "Enhanced prompt" assert "cleaned" in result.model_dump() @pytest.mark.asyncio async def test_response_format_validation( self, mock_system_prompt, mock_llm_response ): """Test that response format matches MCP protocol""" with patch("builtins.open", mock_open(read_data=mock_system_prompt)): with patch("tools.cleaner.LLMClient") as mock_client_class: mock_client = AsyncMock() mock_client.chat_completions.return_value = mock_llm_response mock_client_class.return_value = mock_client from main import clean_prompt_tool result = await clean_prompt_tool(raw_prompt="test prompt") # Verify response is a CleanPromptOutput object assert hasattr(result, "cleaned") assert hasattr(result, "notes") assert hasattr(result, "open_questions") assert hasattr(result, "risks") assert hasattr(result, "unchanged") assert hasattr(result, "quality") # Verify it can be serialized to JSON json_data = result.model_dump() assert isinstance(json_data, dict) assert "cleaned" in json_data