python-sdk/pkg/hanzo/tests/test_batch_orchestrator.py

"""Test suite for batch orchestrator with parallel agent execution."""

import os
import shutil
import asyncio
import tempfile
from pathlib import Path
from datetime import datetime
from unittest.mock import Mock, AsyncMock, MagicMock, patch

import pytest
from hanzo.batch_orchestrator import (
    BatchTask,
    BatchConfig,
    BatchOrchestrator,
    MetaAIOrchestrator,
)


class TestBatchConfig:
    """Test batch configuration parsing."""

    def test_default_config(self):
        """Test default configuration values."""
        config = BatchConfig()
        assert config.batch_size == 5
        assert config.agent_model == "claude-3-5-sonnet-20241022"
        assert config.use_worktrees == False

    def test_parse_simple_batch(self):
        """Test parsing simple batch command."""
        config = BatchConfig.from_command("batch:10 add copyright to files")
        assert config.batch_size == 10
        assert config.agent_model == "claude-3-5-sonnet-20241022"  # Defaults to Claude
        assert config.operation == "add copyright to files"

    def test_parse_with_agent(self):
        """Test parsing with explicit agent."""
        config = BatchConfig.from_command("batch:5 agent:codex fix typing")
        assert config.batch_size == 5
        assert config.agent_model == "gpt-4-turbo"
        assert config.operation == "fix typing"

    def test_parse_with_worktree(self):
        """Test parsing with worktree option."""
        config = BatchConfig.from_command("batch:3 worktree:true agent:gemini add docs")
        assert config.batch_size == 3
        assert config.use_worktrees == True
        assert config.agent_model == "gemini-1.5-pro"
        assert config.operation == "add docs"

    def test_parse_with_files_pattern(self):
        """Test parsing with file pattern."""
        config = BatchConfig.from_command("batch:5 files:*.py add type hints")
        assert config.batch_size == 5
        assert config.target_pattern == "*.py"
        assert config.operation == "add type hints"

    def test_agent_aliases(self):
        """Test agent name aliases."""
        # Test Claude Code alias
        config = BatchConfig.from_command("batch:5 agent:cc refactor")
        assert config.agent_model == "claude-3-5-sonnet-20241022"

        # Test other aliases
        config = BatchConfig.from_command("batch:5 agent:llama optimize")
        assert config.agent_model == "ollama/llama-3.2-3b"

        config = BatchConfig.from_command("batch:5 agent:deepseek analyze")
        assert config.agent_model == "deepseek-coder-v2"


class TestBatchTask:
    """Test batch task functionality."""

    def test_task_creation(self):
        """Test task creation."""
        task = BatchTask(
            id="task_001",
            description="Add copyright",
            file_path=Path("test.py"),
            agent_model="claude-3-5-sonnet-20241022",
        )
        assert task.id == "task_001"
        assert task.status == "pending"
        assert task.result is None
        assert task.error is None

    def test_task_duration(self):
        """Test task duration calculation."""
        task = BatchTask(id="test", description="test")
        assert task.duration() is None

        task.start_time = datetime(2024, 1, 1, 10, 0, 0)
        task.end_time = datetime(2024, 1, 1, 10, 0, 5)
        assert task.duration() == 5.0


class TestBatchOrchestrator:
    """Test batch orchestrator functionality."""

    @pytest.fixture
    def mock_mcp_client(self):
        """Create mock MCP client."""
        client = MagicMock()
        client.call_tool = AsyncMock()
        return client

    @pytest.fixture
    def mock_hanzo_client(self):
        """Create mock Hanzo client."""
        client = MagicMock()
        client.chat.completions.create = AsyncMock()
        return client

    @pytest.fixture
    def orchestrator(self, mock_mcp_client, mock_hanzo_client):
        """Create batch orchestrator instance."""
        return BatchOrchestrator(
            mcp_client=mock_mcp_client,
            hanzo_client=mock_hanzo_client,
        )

    @pytest.mark.asyncio
    async def test_find_target_files_with_mcp(self, orchestrator, mock_mcp_client):
        """Test finding files with MCP client."""
        mock_mcp_client.call_tool.return_value = "file1.py\nfile2.py\nfile3.py"

        files = await orchestrator._find_target_files("*.py")

        assert len(files) == 3
        assert Path("file1.py") in files
        mock_mcp_client.call_tool.assert_called_once_with("find", {"pattern": "*.py"})

    @pytest.mark.asyncio
    async def test_execute_agent_task_with_mcp(self, orchestrator, mock_mcp_client):
        """Test executing task with MCP agent."""
        task = BatchTask(
            id="test_001",
            description="Add copyright",
            file_path=Path("test.py"),
        )
        config = BatchConfig(operation="add copyright")

        mock_mcp_client.call_tool.return_value = "Copyright added successfully"

        await orchestrator._execute_agent_task(task, config)

        assert task.status == "completed"
        assert task.result == "Copyright added successfully"
        assert task.duration() is not None

        mock_mcp_client.call_tool.assert_called_once()
        call_args = mock_mcp_client.call_tool.call_args[0]
        assert call_args[0] == "agent"
        assert "add copyright" in call_args[1]["prompt"]

    @pytest.mark.asyncio
    async def test_execute_agent_task_with_hanzo(self, orchestrator, mock_hanzo_client):
        """Test executing task with Hanzo client."""
        orchestrator.mcp_client = None  # Disable MCP to use Hanzo

        task = BatchTask(
            id="test_002",
            description="Fix typing",
        )
        config = BatchConfig(operation="fix typing", agent_model="gpt-4-turbo")

        # Mock Hanzo response
        mock_response = MagicMock()
        mock_response.choices = [MagicMock()]
        mock_response.choices[0].message.content = "Typing fixed"
        mock_hanzo_client.chat.completions.create.return_value = mock_response

        await orchestrator._execute_agent_task(task, config)

        assert task.status == "completed"
        assert task.result == "Typing fixed"

        mock_hanzo_client.chat.completions.create.assert_called_once()
        call_kwargs = mock_hanzo_client.chat.completions.create.call_args[1]
        assert call_kwargs["model"] == "gpt-4-turbo"

    @pytest.mark.asyncio
    async def test_execute_batch_simple(self, orchestrator, mock_mcp_client):
        """Test executing simple batch operation."""
        # Mock file finding
        mock_mcp_client.call_tool.side_effect = [
            "file1.py\nfile2.py",  # Find files
            "Task 1 complete",  # Agent task 1
            "Task 2 complete",  # Agent task 2
        ]

        summary = await orchestrator.execute_batch("batch:2 add copyright")

        assert summary["total_tasks"] == 2
        assert summary["completed"] == 2
        assert summary["failed"] == 0
        assert summary["batch_size"] == 2
        assert summary["agent_model"] == "claude-3-5-sonnet-20241022"

    @pytest.mark.asyncio
    async def test_parallel_execution(self, orchestrator, mock_mcp_client):
        """Test parallel task execution."""

        # Create delay to test parallelism
        async def delayed_response(*args, **kwargs):
            await asyncio.sleep(0.1)
            return f"Task complete"

        mock_mcp_client.call_tool.side_effect = [
            "file1.py\nfile2.py\nfile3.py\nfile4.py",  # Find files
        ] + [
            delayed_response
        ] * 4  # 4 agent tasks

        import time

        start = time.time()

        summary = await orchestrator.execute_batch("batch:2 process files")

        duration = time.time() - start

        # With batch size 2, should take ~0.2s (2 batches of 0.1s each)
        # Not 0.4s if sequential
        assert duration < 0.3
        assert summary["completed"] == 4
        assert summary["batch_size"] == 2

    @pytest.mark.asyncio
    async def test_task_failure_handling(self, orchestrator, mock_mcp_client):
        """Test handling of failed tasks."""
        mock_mcp_client.call_tool.side_effect = [
            "file1.py\nfile2.py",  # Find files
            "Success",  # Task 1 succeeds
            Exception("API error"),  # Task 2 fails
        ]

        summary = await orchestrator.execute_batch("batch:2 process")

        assert summary["completed"] == 1
        assert summary["failed"] == 1
        assert len(orchestrator.completed_tasks) == 1
        assert len(orchestrator.failed_tasks) == 1
        assert orchestrator.failed_tasks[0].error == "API error"


class TestGitWorktreeIntegration:
    """Test git worktree integration."""

    @pytest.fixture
    def temp_git_repo(self):
        """Create temporary git repository."""
        temp_dir = tempfile.mkdtemp()
        os.chdir(temp_dir)

        # Initialize git repo
        import subprocess

        subprocess.run(["git", "init"], capture_output=True)
        subprocess.run(
            ["git", "config", "user.email", "[email protected]"], capture_output=True
        )
        subprocess.run(["git", "config", "user.name", "Test User"], capture_output=True)

        # Create initial commit
        Path("test.txt").write_text("initial content")
        subprocess.run(["git", "add", "."], capture_output=True)
        subprocess.run(["git", "commit", "-m", "Initial commit"], capture_output=True)

        yield temp_dir

        # Cleanup
        os.chdir("/")
        shutil.rmtree(temp_dir, ignore_errors=True)

    @pytest.mark.asyncio
    async def test_worktree_setup(self, temp_git_repo):
        """Test worktree setup and cleanup."""
        orchestrator = BatchOrchestrator()
        config = BatchConfig(use_worktrees=True)

        # Setup worktree
        worktree_path = await orchestrator._setup_worktree("test_task", config)

        assert worktree_path is not None
        assert worktree_path.exists()
        assert (worktree_path / "test.txt").exists()

        # Cleanup worktree
        await orchestrator._cleanup_worktree("test_task")

        assert not worktree_path.exists()
        assert "test_task" not in orchestrator._worktrees

    @pytest.mark.asyncio
    async def test_worktree_changes_merge(self, temp_git_repo):
        """Test merging changes from worktree."""
        orchestrator = BatchOrchestrator()
        config = BatchConfig(use_worktrees=True)

        # Setup worktree
        worktree_path = await orchestrator._setup_worktree("test_task", config)

        # Make changes in worktree
        (worktree_path / "new_file.txt").write_text("new content")

        # Merge changes
        await orchestrator._merge_worktree_changes("test_task", worktree_path)

        # Cleanup
        await orchestrator._cleanup_worktree("test_task")

        # Check changes are in main branch
        assert Path("new_file.txt").exists()
        assert Path("new_file.txt").read_text() == "new content"


class TestMetaAIOrchestrator:
    """Test meta AI orchestrator."""

    @pytest.fixture
    def meta_orchestrator(self):
        """Create meta AI orchestrator."""
        mock_mcp = MagicMock()
        mock_mcp.call_tool = AsyncMock()
        return MetaAIOrchestrator(mcp_client=mock_mcp)

    @pytest.mark.asyncio
    async def test_parse_batch_command(self, meta_orchestrator):
        """Test parsing batch commands."""
        # Mock file finding to avoid filesystem scan
        meta_orchestrator.mcp_client.call_tool.return_value = "file1.py\nfile2.py"

        result = await meta_orchestrator.parse_and_execute(
            "batch:5 files:*.py add copyright"
        )

        # Should recognize as batch command
        assert "total_tasks" in result or "error" not in result

    @pytest.mark.asyncio
    async def test_natural_language_to_batch(self, meta_orchestrator):
        """Test converting natural language to batch."""
        # Mock intent analysis
        meta_orchestrator.mcp_client.call_tool.return_value = """
        {
            "type": "batch_operation",
            "operation": "add copyright headers",
            "model": "claude-3-5-sonnet-20241022",
            "pattern": "*.py",
            "batch_size": 10
        }
        """

        intent = await meta_orchestrator._analyze_intent(
            "Add copyright headers to all Python files"
        )

        assert intent["type"] == "batch_operation"

        # Build batch command
        batch_cmd = meta_orchestrator._build_batch_command(intent)
        assert "batch:10" in batch_cmd
        assert "claude" in batch_cmd

    @pytest.mark.asyncio
    async def test_single_task_execution(self, meta_orchestrator):
        """Test single task execution."""
        meta_orchestrator.mcp_client.call_tool.return_value = """
        {
            "type": "single_task",
            "operation": "explain this code",
            "model": "claude-3-5-sonnet-20241022"
        }
        """

        result = await meta_orchestrator.parse_and_execute("explain this code")

        # Should execute as single task
        assert "task_id" in result or "error" not in result


@pytest.mark.integration
class TestBatchOrchestratorIntegration:
    """Integration tests for batch orchestrator."""

    @pytest.mark.asyncio
    async def test_real_batch_execution(self):
        """Test real batch execution with files."""
        # This test requires actual MCP client
        pytest.skip("Requires MCP server running")

        from hanzo_mcp import MCPClient

        client = MCPClient()
        orchestrator = BatchOrchestrator(mcp_client=client)

        # Create test files
        test_dir = Path("test_batch")
        test_dir.mkdir(exist_ok=True)

        for i in range(3):
            (test_dir / f"file{i}.txt").write_text(f"Content {i}")

        try:
            # Execute batch
            summary = await orchestrator.execute_batch(
                "batch:3 files:test_batch/*.txt add header"
            )

            assert summary["completed"] == 3
            assert summary["failed"] == 0

        finally:
            # Cleanup
            shutil.rmtree(test_dir, ignore_errors=True)