mcp-forge/tests/server/tools/test_build_environment.py

"""Tests for Build Custom Environment Tool."""

import pytest
from unittest.mock import Mock, AsyncMock
from mcp.types import Tool, TextContent
import json

from mcp_forge.server.tools.build_environment import BuildEnvironmentTool
from mcp_forge.builder.environment_builder import BuildResult


@pytest.fixture
def mock_environment_builder():
    """Mock environment builder."""
    builder = Mock()
    builder.build_environment = AsyncMock(return_value=BuildResult(
        success=True,
        image_name="mcp-forge/custom:test-env",
        image_id="sha256:abc123",
        build_time=45.2,
        size_bytes=524288000,  # 500MB
        installed_packages=["numpy==1.24.0", "pandas==2.0.0"],
        cache_hit=False
    ))
    return builder


@pytest.fixture
def mock_audit_logger():
    """Mock audit logger."""
    logger = Mock()
    logger.log_environment_build = Mock()
    return logger


@pytest.fixture
def build_environment_tool(mock_environment_builder, mock_audit_logger):
    """Create BuildEnvironmentTool instance with mocked dependencies."""
    return BuildEnvironmentTool(
        environment_builder=mock_environment_builder,
        audit_logger=mock_audit_logger
    )


@pytest.mark.asyncio
async def test_get_tool_definition(build_environment_tool):
    """Test that tool definition matches MCP spec."""
    definition = build_environment_tool.get_tool_definition()
    
    assert isinstance(definition, Tool)
    assert definition.name == "build_custom_environment"
    assert definition.description is not None
    assert "build" in definition.description.lower()
    
    # Verify required schema properties
    schema = definition.inputSchema
    assert schema["type"] == "object"
    assert "name" in schema["properties"]
    assert "packages" in schema["properties"]
    assert "base_image" in schema["properties"]
    assert "python_version" in schema["properties"]
    assert "description" in schema["properties"]
    assert "name" in schema["required"]
    assert "packages" in schema["required"]


@pytest.mark.asyncio
async def test_build_with_valid_packages(
    build_environment_tool,
    mock_environment_builder,
    mock_audit_logger
):
    """Test building environment with valid packages."""
    arguments = {
        "name": "test-ml-env",
        "packages": ["numpy>=1.24.0", "pandas>=2.0.0"]
    }
    
    result = await build_environment_tool.execute(arguments)
    
    assert isinstance(result, list)
    assert len(result) == 1
    assert isinstance(result[0], TextContent)
    
    # Parse JSON response
    response = json.loads(result[0].text)
    assert response["success"] is True
    assert response["image_name"] == "mcp-forge/custom:test-env"
    assert response["build_time"] == 45.2
    assert len(response["installed_packages"]) == 2
    
    # Verify builder was called
    mock_environment_builder.build_environment.assert_called_once()
    
    # Verify audit log was called
    mock_audit_logger.log_environment_build.assert_called_once()


@pytest.mark.asyncio
async def test_build_with_base_image(
    build_environment_tool,
    mock_environment_builder
):
    """Test building with custom base image."""
    arguments = {
        "name": "test-env",
        "packages": ["requests"],
        "base_image": "python:3.12"
    }
    
    result = await build_environment_tool.execute(arguments)
    
    # Parse JSON response
    response = json.loads(result[0].text)
    assert response["success"] is True
    
    # Verify base_image was passed to builder
    call_args = mock_environment_builder.build_environment.call_args
    assert call_args.kwargs["base_image"] == "python:3.12"


@pytest.mark.asyncio
async def test_build_with_python_version(
    build_environment_tool,
    mock_environment_builder
):
    """Test building with specific Python version."""
    arguments = {
        "name": "test-env",
        "packages": ["numpy"],
        "python_version": "3.12"
    }
    
    result = await build_environment_tool.execute(arguments)
    
    # Parse JSON response
    response = json.loads(result[0].text)
    assert response["success"] is True
    
    # Verify python_version was passed to builder
    call_args = mock_environment_builder.build_environment.call_args
    assert call_args.kwargs["python_version"] == "3.12"


@pytest.mark.asyncio
async def test_build_with_description(
    build_environment_tool,
    mock_environment_builder
):
    """Test building with description."""
    arguments = {
        "name": "test-env",
        "packages": ["pandas"],
        "description": "Environment for data analysis"
    }
    
    result = await build_environment_tool.execute(arguments)
    
    # Parse JSON response
    response = json.loads(result[0].text)
    assert response["success"] is True


@pytest.mark.asyncio
async def test_validate_arguments_missing_name(build_environment_tool):
    """Test that validation fails when name is missing."""
    arguments = {
        "packages": ["numpy"]
    }
    
    with pytest.raises((ValueError, KeyError)):
        await build_environment_tool.execute(arguments)


@pytest.mark.asyncio
async def test_validate_arguments_missing_packages(build_environment_tool):
    """Test that validation fails when packages is missing."""
    arguments = {
        "name": "test-env"
    }
    
    with pytest.raises((ValueError, KeyError)):
        await build_environment_tool.execute(arguments)


@pytest.mark.asyncio
async def test_validate_arguments_invalid_name(build_environment_tool):
    """Test that validation fails for invalid environment name."""
    arguments = {
        "name": "invalid name with spaces",
        "packages": ["numpy"]
    }
    
    with pytest.raises(ValueError, match="name"):
        await build_environment_tool.execute(arguments)


@pytest.mark.asyncio
async def test_validate_arguments_invalid_packages_type(build_environment_tool):
    """Test that validation fails when packages is not a list."""
    arguments = {
        "name": "test-env",
        "packages": "not-a-list"
    }
    
    with pytest.raises(ValueError, match="packages"):
        await build_environment_tool.execute(arguments)


@pytest.mark.asyncio
async def test_validate_arguments_empty_packages(build_environment_tool):
    """Test that validation fails when packages list is empty."""
    arguments = {
        "name": "test-env",
        "packages": []
    }
    
    with pytest.raises(ValueError, match="packages"):
        await build_environment_tool.execute(arguments)


@pytest.mark.asyncio
async def test_build_error_handling(
    build_environment_tool,
    mock_environment_builder
):
    """Test that build errors are handled gracefully."""
    mock_environment_builder.build_environment = AsyncMock(return_value=BuildResult(
        success=False,
        image_name="",
        image_id="",
        build_time=5.0,
        size_bytes=0,
        installed_packages=[],
        cache_hit=False,
        error="Package 'invalid-pkg' not found"
    ))
    
    arguments = {
        "name": "test-env",
        "packages": ["invalid-pkg"]
    }
    
    result = await build_environment_tool.execute(arguments)
    
    # Parse JSON response
    response = json.loads(result[0].text)
    assert response["success"] is False
    assert "error" in response
    assert "invalid-pkg" in response["error"]