CodeMerge

#!/usr/bin/env python3
"""
Tests for the codemerge MCP server.
"""

import asyncio
import os
import tempfile
import unittest
from unittest.mock import Mock, patch, AsyncMock
from typing import Dict, Any

# Import the functions we want to test
from codemerge import edit_snippet, list_tools, call_tool


class TestCodemerge(unittest.TestCase):
    """Test cases for the codemerge MCP server."""
    
    def setUp(self):
        """Set up test fixtures."""
        self.sample_original_code = """def hello_world():
    print("Hello, World!")
    return "success" """
        
        self.sample_edit_snippet = """def hello_world():
    print("Hello, Universe!")
    # ... existing code ...
    return "success" """
        
        self.expected_edited_code = """def hello_world():
    print("Hello, Universe!")
    return "success" """
        
        # Mock ollama response
        self.mock_ollama_response = {
            'message': {
                'content': self.expected_edited_code + ' \n</code>'
            }
        }
    
    def test_list_tools(self):
        """Test that list_tools returns the correct tool definition."""
        async def run_test():
            tools = await list_tools()
            self.assertEqual(len(tools), 1)
            
            tool = tools[0]
            self.assertEqual(tool.name, "edit_snippet")
            self.assertIn("Apply an edit to a code snippet", tool.description)
            
            # Check required parameters
            required = tool.inputSchema["required"]
            self.assertIn("original_code", required)
            self.assertIn("edit_snippet", required)
            self.assertNotIn("file_path", required)  # Optional parameter
            
            # Check properties
            properties = tool.inputSchema["properties"]
            self.assertIn("original_code", properties)
            self.assertIn("edit_snippet", properties)
            self.assertIn("file_path", properties)
        
        asyncio.run(run_test())
    
    def test_call_tool_edit_snippet(self):
        """Test call_tool with edit_snippet."""
        async def run_test():
            with patch('codemerge.edit_snippet') as mock_edit:
                mock_edit.return_value = [{"type": "text", "text": "test result"}]
                
                arguments = {
                    "original_code": self.sample_original_code,
                    "edit_snippet": self.sample_edit_snippet
                }
                
                result = await call_tool("edit_snippet", arguments)
                
                mock_edit.assert_called_once_with(arguments)
                self.assertEqual(result, [{"type": "text", "text": "test result"}])
        
        asyncio.run(run_test())
    
    def test_call_tool_unknown_tool(self):
        """Test call_tool with unknown tool name."""
        async def run_test():
            with self.assertRaises(ValueError) as context:
                await call_tool("unknown_tool", {})
            
            self.assertIn("Unknown tool: unknown_tool", str(context.exception))
        
        asyncio.run(run_test())


class TestEditSnippet(unittest.TestCase):
    """Test cases for the edit_snippet function."""
    
    def setUp(self):
        """Set up test fixtures."""
        self.sample_original_code = """def hello_world():
    print("Hello, World!")
    return "success" """
        
        self.sample_edit_snippet = """def hello_world():
    print("Hello, Universe!")
    # ... existing code ...
    return "success" """
        
        self.expected_edited_code = """def hello_world():
    print("Hello, Universe!")
    return "success" """
        
        # Mock ollama response
        self.mock_ollama_response = {
            'message': {
                'content': self.expected_edited_code + ' \n</code>'
            }
        }
    
    @patch('codemerge.ollama.chat')
    def test_edit_snippet_without_file_path(self, mock_ollama):
        """Test edit_snippet without file_path parameter."""
        async def run_test():
            mock_ollama.return_value = self.mock_ollama_response
            
            arguments = {
                "original_code": self.sample_original_code,
                "edit_snippet": self.sample_edit_snippet
            }
            
            result = await edit_snippet(arguments)
            
            # Verify ollama was called
            mock_ollama.assert_called_once()
            
            # Verify result
            self.assertEqual(len(result), 1)
            self.assertEqual(result[0]["type"], "text")
            self.assertEqual(result[0]["text"].strip(), self.expected_edited_code.strip())
        
        asyncio.run(run_test())
    
    @patch('codemerge.ollama.chat')
    def test_edit_snippet_with_file_path_success(self, mock_ollama):
        """Test edit_snippet with file_path parameter - successful replacement."""
        async def run_test():
            mock_ollama.return_value = self.mock_ollama_response
            
            # Create a temporary file with the original code
            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as temp_file:
                temp_file.write(f"# Some header\n{self.sample_original_code}\n# Some footer")
                temp_file_path = temp_file.name
            
            try:
                arguments = {
                    "original_code": self.sample_original_code,
                    "edit_snippet": self.sample_edit_snippet,
                    "file_path": temp_file_path
                }
                
                result = await edit_snippet(arguments)
                
                # Verify result
                self.assertEqual(len(result), 1)
                self.assertEqual(result[0]["type"], "text")
                self.assertIn(self.expected_edited_code, result[0]["text"])
                self.assertIn("File operation: SUCCESS", result[0]["text"])
                
                # Verify file was updated
                with open(temp_file_path, 'r') as f:
                    updated_content = f.read()
                
                self.assertIn(self.expected_edited_code, updated_content)
                self.assertNotIn(self.sample_original_code, updated_content)
                
            finally:
                # Clean up
                os.unlink(temp_file_path)
        
        asyncio.run(run_test())
    
    @patch('codemerge.ollama.chat')
    def test_edit_snippet_with_file_path_code_not_found(self, mock_ollama):
        """Test edit_snippet with file_path when original code is not found."""
        async def run_test():
            mock_ollama.return_value = self.mock_ollama_response
            
            # Create a temporary file without the original code
            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as temp_file:
                temp_file.write("# Some other code\nprint('different code')")
                temp_file_path = temp_file.name
            
            try:
                arguments = {
                    "original_code": self.sample_original_code,
                    "edit_snippet": self.sample_edit_snippet,
                    "file_path": temp_file_path
                }
                
                result = await edit_snippet(arguments)
                
                # Verify result
                self.assertEqual(len(result), 1)
                self.assertEqual(result[0]["type"], "text")
                self.assertIn(self.expected_edited_code, result[0]["text"])
                self.assertIn("File operation: FAILED - Original code not found", result[0]["text"])
                
            finally:
                # Clean up
                os.unlink(temp_file_path)
        
        asyncio.run(run_test())
    
    @patch('codemerge.ollama.chat')
    def test_edit_snippet_with_nonexistent_file(self, mock_ollama):
        """Test edit_snippet with non-existent file path."""
        async def run_test():
            mock_ollama.return_value = self.mock_ollama_response
            
            arguments = {
                "original_code": self.sample_original_code,
                "edit_snippet": self.sample_edit_snippet,
                "file_path": "/nonexistent/path/file.py"
            }
            
            result = await edit_snippet(arguments)
            
            # Verify result
            self.assertEqual(len(result), 1)
            self.assertEqual(result[0]["type"], "text")
            self.assertIn(self.expected_edited_code, result[0]["text"])
            self.assertIn("File operation: FAILED - File", result[0]["text"])
            self.assertIn("does not exist", result[0]["text"])
        
        asyncio.run(run_test())
    
    @patch('codemerge.ollama.chat')
    def test_edit_snippet_ollama_error(self, mock_ollama):
        """Test edit_snippet when ollama raises an exception."""
        async def run_test():
            mock_ollama.side_effect = Exception("Ollama connection error")
            
            arguments = {
                "original_code": self.sample_original_code,
                "edit_snippet": self.sample_edit_snippet
            }
            
            result = await edit_snippet(arguments)
            
            # Verify error handling
            self.assertEqual(len(result), 1)
            self.assertEqual(result[0]["type"], "text")
            self.assertIn("Error applying edit: Ollama connection error", result[0]["text"])
        
        asyncio.run(run_test())
    
    @patch('codemerge.ollama.chat')
    def test_edit_snippet_file_permission_error(self, mock_ollama):
        """Test edit_snippet with file permission error."""
        async def run_test():
            mock_ollama.return_value = self.mock_ollama_response
            
            # Create a temporary file and make it read-only
            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as temp_file:
                temp_file.write(f"# Some header\n{self.sample_original_code}\n# Some footer")
                temp_file_path = temp_file.name
            
            try:
                # Make file read-only
                os.chmod(temp_file_path, 0o444)
                
                arguments = {
                    "original_code": self.sample_original_code,
                    "edit_snippet": self.sample_edit_snippet,
                    "file_path": temp_file_path
                }
                
                result = await edit_snippet(arguments)
                
                # Verify result
                self.assertEqual(len(result), 1)
                self.assertEqual(result[0]["type"], "text")
                self.assertIn(self.expected_edited_code, result[0]["text"])
                self.assertIn("File operation: FAILED - Error:", result[0]["text"])
                
            finally:
                # Clean up - restore permissions and delete
                os.chmod(temp_file_path, 0o644)
                os.unlink(temp_file_path)
        
        asyncio.run(run_test())
    
    @patch('codemerge.ollama.chat')
    def test_edit_snippet_malformed_response(self, mock_ollama):
        """Test edit_snippet with malformed ollama response."""
        async def run_test():
            # Mock response without proper code tags
            mock_ollama.return_value = {
                'message': {
                    'content': 'Some response without code tags'
                }
            }
            
            arguments = {
                "original_code": self.sample_original_code,
                "edit_snippet": self.sample_edit_snippet
            }
            
            result = await edit_snippet(arguments)
            
            # Verify result uses fallback content
            self.assertEqual(len(result), 1)
            self.assertEqual(result[0]["type"], "text")
            self.assertEqual(result[0]["text"], "Some response without code tags")
        
        asyncio.run(run_test())


if __name__ == '__main__':
    unittest.main()