python-sdk/pkg/hanzo-repl/src/hanzo_repl/tests.py

"""Test suite for MCP tools in REPL context."""

import tempfile
from pathlib import Path
from typing import Any

from hanzo_mcp.server import HanzoMCPServer
from rich.console import Console
from rich.table import Table


async def run_tool_tests(
    console: Console, mcp_server: HanzoMCPServer, tool_executor: Any
):
    """Run comprehensive tests for MCP tools."""

    console.print("\n[bold cyan]Running MCP Tool Tests[/bold cyan]\n")

    test_results = []

    # Test categories
    test_suites = [
        ("File Operations", test_file_operations),
        ("Search Operations", test_search_operations),
        ("Shell Commands", test_shell_commands),
        ("LLM Integration", test_llm_integration),
        ("Agent Delegation", test_agent_delegation),
    ]

    for suite_name, test_func in test_suites:
        console.print(f"\n[bold]{suite_name}[/bold]")
        try:
            results = await test_func(mcp_server, tool_executor)
            for test_name, success, message in results:
                test_results.append((suite_name, test_name, success, message))
                status = "[green]✓[/green]" if success else "[red]✗[/red]"
                console.print(f"  {status} {test_name}: {message}")
        except Exception as e:
            console.print(f"  [red]Suite failed: {e}[/red]")

    # Summary
    console.print("\n[bold]Test Summary[/bold]")
    table = Table()
    table.add_column("Suite", style="cyan")
    table.add_column("Test", style="white")
    table.add_column("Status", style="green")

    passed = sum(1 for _, _, success, _ in test_results if success)
    total = len(test_results)

    for suite, test, success, _ in test_results:
        status = "PASS" if success else "FAIL"
        style = "green" if success else "red"
        table.add_row(suite, test, f"[{style}]{status}[/{style}]")

    console.print(table)
    console.print(f"\nTotal: {passed}/{total} passed ({passed / total * 100:.1f}%)")


async def test_file_operations(
    mcp_server: HanzoMCPServer,
    tool_executor: Any,  # noqa: ARG001
) -> list:
    """Test file operations."""
    results = []

    with tempfile.TemporaryDirectory() as tmpdir:
        test_file = Path(tmpdir) / "test.txt"

        # Test 1: Write file
        try:
            tool = mcp_server.tools.get("write_file")
            await tool.execute(file_path=str(test_file), content="Hello, MCP!")
            results.append(("Write file", True, "File created successfully"))
        except Exception as e:
            results.append(("Write file", False, str(e)))

        # Test 2: Read file
        try:
            tool = mcp_server.tools.get("read_file")
            content = await tool.execute(file_path=str(test_file))
            success = "Hello, MCP!" in content
            results.append(("Read file", success, f"Content: {content[:50]}..."))
        except Exception as e:
            results.append(("Read file", False, str(e)))

        # Test 3: Edit file
        try:
            tool = mcp_server.tools.get("edit_file")
            await tool.execute(
                file_path=str(test_file),
                old_string="Hello, MCP!",
                new_string="Hello, Hanzo MCP!",
            )
            # Verify edit
            read_tool = mcp_server.tools.get("read_file")
            content = await read_tool.execute(file_path=str(test_file))
            success = "Hello, Hanzo MCP!" in content
            results.append(("Edit file", success, "File edited successfully"))
        except Exception as e:
            results.append(("Edit file", False, str(e)))

    return results


async def test_search_operations(
    mcp_server: HanzoMCPServer,
    tool_executor: Any,  # noqa: ARG001
) -> list:
    """Test search operations."""
    results = []

    # Test 1: Grep search
    try:
        tool = mcp_server.tools.get("grep")
        result = await tool.execute(pattern="def ", path=".", include="*.py")
        success = isinstance(result, (list, str))
        results.append(
            (
                "Grep search",
                success,
                f"Found {len(result) if isinstance(result, list) else 1} matches",
            )
        )
    except Exception as e:
        results.append(("Grep search", False, str(e)))

    # Test 2: File search
    try:
        tool = mcp_server.tools.get("search")
        result = await tool.execute(query="class", path=".", max_results=5)
        success = isinstance(result, (list, dict, str))
        results.append(("File search", success, "Search completed"))
    except Exception as e:
        results.append(("File search", False, str(e)))

    return results


async def test_shell_commands(
    mcp_server: HanzoMCPServer,
    tool_executor: Any,  # noqa: ARG001
) -> list:
    """Test shell command execution."""
    results = []

    # Test 1: Simple command
    try:
        tool = mcp_server.tools.get("run_command")
        result = await tool.execute(command="echo 'Hello from MCP'")
        success = "Hello from MCP" in str(result)
        results.append(("Echo command", success, "Command executed"))
    except Exception as e:
        results.append(("Echo command", False, str(e)))

    # Test 2: Python version
    try:
        tool = mcp_server.tools.get("run_command")
        result = await tool.execute(command="python --version")
        success = "Python" in str(result)
        results.append(("Python version", success, str(result)[:50]))
    except Exception as e:
        results.append(("Python version", False, str(e)))

    return results


async def test_llm_integration(
    mcp_server: HanzoMCPServer,  # noqa: ARG001
    tool_executor: Any,
) -> list:
    """Test LLM integration."""
    results = []

    # Test 1: Simple chat
    try:
        response = await tool_executor.execute_with_tools("What is 2+2?")
        success = "4" in response
        results.append(("Simple math", success, "LLM responded correctly"))
    except Exception as e:
        results.append(("Simple math", False, str(e)))

    # Test 2: Tool usage
    try:
        response = await tool_executor.execute_with_tools(
            "Create a file called test_llm.txt with the content 'LLM test'"
        )
        # Check if file was created
        import os

        success = os.path.exists("test_llm.txt")
        if success:
            os.remove("test_llm.txt")  # Cleanup
        results.append(("LLM tool use", success, "LLM used tools correctly"))
    except Exception as e:
        results.append(("LLM tool use", False, str(e)))

    return results


async def test_agent_delegation(
    mcp_server: HanzoMCPServer,
    tool_executor: Any,  # noqa: ARG001
) -> list:
    """Test agent delegation."""
    results = []

    # Test 1: Agent dispatch
    try:
        tool = mcp_server.tools.get("dispatch_agent")
        if tool:
            result = await tool.execute(
                instruction="List the current directory contents"
            )
            success = result is not None
            results.append(("Agent dispatch", success, "Agent completed task"))
        else:
            results.append(("Agent dispatch", False, "Agent tool not available"))
    except Exception as e:
        results.append(("Agent dispatch", False, str(e)))

    return results