mcp-forge/tests/builder/test_image_builder.py

"""Tests for Image Builder module."""

import pytest
from pathlib import Path
from unittest.mock import Mock

from mcp_forge.builder.image_builder import ImageBuilder, BuildResult
from pod_executor.containers.client import PodmanClient
from mcp_forge.config.schema import EnvironmentBuilderConfig
from mcp_forge.security.audit import AuditLogger


@pytest.fixture
def mock_podman_client():
    """Mock PodmanClient."""
    mock = Mock(spec=PodmanClient)
    # Configure nested mocks for images and containers
    mock.images = Mock()
    mock.containers = Mock()
    return mock


@pytest.fixture
def builder_config():
    """Mock EnvironmentBuilderConfig."""
    config = Mock(spec=EnvironmentBuilderConfig)
    config.enabled = True
    config.uv_cache_path = Path("/tmp/uv_cache")
    config.max_packages_per_build = 50
    config.build_timeout = 600
    config.max_build_timeout = 1800
    config.max_image_size = "2g"
    return config


@pytest.fixture
def mock_audit_logger():
    """Mock AuditLogger."""
    return Mock(spec=AuditLogger)


@pytest.fixture
def builder(mock_podman_client, builder_config, mock_audit_logger):
    """ImageBuilder instance."""
    return ImageBuilder(mock_podman_client, builder_config, mock_audit_logger)


@pytest.fixture
def build_context(tmp_path):
    """Create temporary build context."""
    context = tmp_path / "build_context"
    context.mkdir()
    
    # Create Containerfile
    containerfile = context / "Containerfile"
    containerfile.write_text("FROM python:3.11-slim\n")
    
    # Create requirements.txt
    requirements = context / "requirements.txt"
    requirements.write_text("numpy>=1.24.0\npandas\n")
    
    return context


def test_build_image_success(builder, build_context, mock_podman_client):
    """Test successful image build."""
    # Mock successful build
    mock_image = Mock()
    mock_image.id = "sha256:abc123"
    mock_image.attrs = {"Size": 500_000_000}
    
    mock_podman_client.images.build.return_value = (mock_image, [])
    mock_podman_client.images.get.return_value = mock_image  # Mock images.get() call
    
    # Mock pip list output
    mock_container = Mock()
    mock_container.exec_run.return_value = (
        0,
        b'[{"name": "numpy", "version": "1.24.0"}, {"name": "pandas", "version": "2.0.0"}]'
    )
    mock_podman_client.containers.run.return_value = mock_container
    
    packages = ["numpy>=1.24.0", "pandas"]
    result = builder.build_image(
        name="test-env",
        build_context=build_context,
        base_image="python:3.11-slim",
        packages=packages
    )
    
    assert result.success is True
    assert result.image_name == "mcp-forge/custom:test-env"
    assert result.image_id == "sha256:abc123"
    assert result.size_bytes == 500_000_000
    assert len(result.installed_packages) == 2


def test_build_image_with_custom_timeout(builder, build_context, mock_podman_client):
    """Test build with custom timeout."""
    mock_image = Mock()
    mock_image.id = "sha256:abc123"
    mock_image.attrs = {"Size": 100_000_000}
    
    mock_podman_client.images.build.return_value = (mock_image, [])
    mock_podman_client.images.get.return_value = mock_image
    mock_podman_client.containers.run.return_value.exec_run.return_value = (0, b"[]")
    
    builder.build_image(
        name="test-env",
        build_context=build_context,
        base_image="python:3.11-slim",
        packages=[],
        timeout=1200
    )
    
    # Verify timeout was passed to build
    call_kwargs = mock_podman_client.images.build.call_args[1]
    assert call_kwargs["timeout"] == 1200


def test_build_image_validates_timeout_against_max(builder, build_context, builder_config):
    """Test that build validates timeout against maximum."""
    builder_config.max_build_timeout = 1800
    
    with pytest.raises(ValueError, match="Timeout 3600 exceeds maximum"):
        builder.build_image(
            name="test-env",
            build_context=build_context,
            base_image="python:3.11-slim",
            packages=[],
            timeout=3600
        )


def test_build_image_uses_default_timeout(builder, build_context, builder_config, mock_podman_client):
    """Test that build uses config default timeout when not specified."""
    builder_config.build_timeout = 600
    
    mock_image = Mock()
    mock_image.id = "sha256:abc123"
    mock_image.attrs = {"Size": 100_000_000}
    
    mock_podman_client.images.build.return_value = (mock_image, [])
    mock_podman_client.images.get.return_value = mock_image
    mock_podman_client.containers.run.return_value.exec_run.return_value = (0, b"[]")
    
    builder.build_image(
        name="test-env",
        build_context=build_context,
        base_image="python:3.11-slim",
        packages=[]
    )
    
    call_kwargs = mock_podman_client.images.build.call_args[1]
    assert call_kwargs["timeout"] == 600


def test_generate_tag_valid_name(builder):
    """Test tag generation with valid name."""
    tag = builder.generate_tag("my-env-123")
    assert tag == "mcp-forge/custom:my-env-123"


def test_generate_tag_invalid_name_raises_error(builder):
    """Test that invalid names raise ValueError."""
    with pytest.raises(ValueError, match="alphanumeric"):
        builder.generate_tag("my env")
    
    with pytest.raises(ValueError, match="alphanumeric"):
        builder.generate_tag("my_env!")


def test_validate_image_size_within_limit(builder, mock_podman_client, builder_config):
    """Test image size validation passes when within limit."""
    builder_config.max_image_size = "2g"
    
    mock_image = Mock()
    mock_image.attrs = {"Size": 1_000_000_000}  # 1GB
    mock_podman_client.images.get.return_value = mock_image
    
    size = builder.validate_image_size("sha256:abc123")
    
    assert size == 1_000_000_000


def test_validate_image_size_exceeds_limit(builder, mock_podman_client, builder_config):
    """Test image size validation fails when exceeds limit."""
    builder_config.max_image_size = "1g"
    
    mock_image = Mock()
    mock_image.attrs = {"Size": 2_000_000_000}  # 2GB
    mock_podman_client.images.get.return_value = mock_image
    
    with pytest.raises(ValueError, match="exceeds maximum"):
        builder.validate_image_size("sha256:abc123")


def test_extract_installed_packages(builder, mock_podman_client):
    """Test extracting installed packages from image."""
    mock_container = Mock()
    mock_container.exec_run.return_value = (
        0,
        b'[{"name": "numpy", "version": "1.24.0"}, {"name": "pandas", "version": "2.0.0"}]'
    )
    mock_podman_client.containers.run.return_value = mock_container
    
    packages = builder.extract_installed_packages("sha256:abc123")
    
    assert len(packages) == 2
    assert "numpy==1.24.0" in packages
    assert "pandas==2.0.0" in packages


def test_extract_installed_packages_handles_error(builder, mock_podman_client):
    """Test that package extraction handles errors gracefully."""
    mock_container = Mock()
    mock_container.exec_run.return_value = (1, b"Error")
    mock_podman_client.containers.run.return_value = mock_container
    
    packages = builder.extract_installed_packages("sha256:abc123")
    
    assert packages == []


def test_calculate_cache_hash(builder):
    """Test cache hash calculation."""
    packages = ["numpy>=1.24.0", "pandas==2.0.0", "requests"]
    
    hash1 = builder.calculate_cache_hash(packages)
    hash2 = builder.calculate_cache_hash(packages)
    
    # Same packages should produce same hash
    assert hash1 == hash2
    
    # Different packages should produce different hash
    different_packages = ["numpy>=1.24.0", "scipy"]
    hash3 = builder.calculate_cache_hash(different_packages)
    assert hash1 != hash3


def test_calculate_cache_hash_order_independent(builder):
    """Test that cache hash is order-independent."""
    packages1 = ["numpy", "pandas", "scipy"]
    packages2 = ["scipy", "numpy", "pandas"]
    
    hash1 = builder.calculate_cache_hash(packages1)
    hash2 = builder.calculate_cache_hash(packages2)
    
    # Order shouldn't matter
    assert hash1 == hash2


def test_build_image_logs_audit_event(builder, build_context, mock_podman_client, mock_audit_logger):
    """Test that build logs audit event."""
    mock_image = Mock()
    mock_image.id = "sha256:abc123"
    mock_image.attrs = {"Size": 100_000_000}
    
    mock_podman_client.images.build.return_value = (mock_image, [])
    mock_podman_client.images.get.return_value = mock_image
    mock_podman_client.containers.run.return_value.exec_run.return_value = (0, b"[]")
    
    builder.build_image(
        name="test-env",
        build_context=build_context,
        base_image="python:3.11-slim",
        packages=["numpy"]
    )
    
    # Verify audit log was called
    assert mock_audit_logger.log.called


def test_build_image_returns_build_time(builder, build_context, mock_podman_client):
    """Test that build result includes build time."""
    mock_image = Mock()
    mock_image.id = "sha256:abc123"
    mock_image.attrs = {"Size": 100_000_000}
    
    mock_podman_client.images.build.return_value = (mock_image, [])
    mock_podman_client.images.get.return_value = mock_image
    mock_podman_client.containers.run.return_value.exec_run.return_value = (0, b"[]")
    
    result = builder.build_image(
        name="test-env",
        build_context=build_context,
        base_image="python:3.11-slim",
        packages=[]
    )
    
    assert result.build_time > 0


def test_build_image_handles_build_failure(builder, build_context, mock_podman_client):
    """Test that build handles Podman build failures."""
    mock_podman_client.images.build.side_effect = Exception("Build failed")
    
    with pytest.raises(RuntimeError, match="Build failed"):
        builder.build_image(
            name="test-env",
            build_context=build_context,
            base_image="python:3.11-slim",
            packages=[]
        )


def test_build_result_to_dict(builder):
    """Test BuildResult serialization."""
    result = BuildResult(
        success=True,
        image_name="mcp-forge/custom:test",
        image_id="sha256:abc123",
        build_time=10.5,
        size_bytes=500_000_000,
        cache_hit=False,
        installed_packages=["numpy==1.24.0", "pandas==2.0.0"]
    )
    
    result_dict = result.to_dict()
    
    assert result_dict["success"] is True
    assert result_dict["image_name"] == "mcp-forge/custom:test"
    assert result_dict["image_id"] == "sha256:abc123"
    assert result_dict["build_time"] == 10.5
    assert result_dict["size_bytes"] == 500_000_000
    assert result_dict["cache_hit"] is False
    assert len(result_dict["installed_packages"]) == 2


def test_build_image_validates_build_context_exists(builder, tmp_path):
    """Test that build validates build context exists."""
    nonexistent = tmp_path / "nonexistent"
    
    with pytest.raises(ValueError, match="Build context does not exist"):
        builder.build_image(
            name="test-env",
            build_context=nonexistent,
            base_image="python:3.11-slim",
            packages=[]
        )


def test_build_image_validates_containerfile_exists(builder, tmp_path):
    """Test that build validates Containerfile exists."""
    context = tmp_path / "context"
    context.mkdir()
    # No Containerfile created
    
    with pytest.raises(ValueError, match="Containerfile not found"):
        builder.build_image(
            name="test-env",
            build_context=context,
            base_image="python:3.11-slim",
            packages=[]
        )