mcp-forge/tests/builder/test_environment_builder.py

"""Tests for environment builder orchestration."""

import pytest
from pathlib import Path
from unittest.mock import Mock, patch
from datetime import datetime, timedelta

from mcp_forge.builder.environment_builder import (
    EnvironmentBuilder,
    BuildRateLimiter,
)
from mcp_forge.builder.package_validator import SecurityError
from mcp_forge.builder.image_builder import BuildResult
from mcp_forge.config.schema import EnvironmentBuilderConfig
from pod_executor.containers.client import PodmanClient
from mcp_forge.security.audit import AuditLogger


@pytest.fixture
def builder_config():
    """Mock environment builder configuration."""
    config = Mock(spec=EnvironmentBuilderConfig)
    config.build_timeout = 600
    config.max_build_timeout = 1800
    config.max_image_size = "2g"
    config.max_packages = 50
    config.uv_cache_path = Path("/tmp/uv-cache")
    config.build_rate_limit = {"requests": 5, "period": 3600}
    config.max_concurrent_builds = 3
    config.base_images = {"python:3.11-slim": True}
    config.templates = {
        "data-science": {
            "packages": ["numpy", "pandas", "matplotlib"],
            "description": "Data science environment"
        }
    }
    # Add package_validation config for PackageValidator
    config.package_validation = Path("/tmp/package-validation.yaml")
    return config


@pytest.fixture
def mock_podman_client():
    """Mock Podman client."""
    return Mock(spec=PodmanClient)


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


@pytest.fixture
def builder(builder_config, mock_podman_client, mock_audit_logger):
    """Environment builder instance."""
    # Mock the sub-components during initialization
    with patch('mcp_forge.builder.environment_builder.PackageValidator'), \
         patch('mcp_forge.builder.environment_builder.UVInstaller'), \
         patch('mcp_forge.builder.environment_builder.ImageBuilder'):
        
        builder = EnvironmentBuilder(
            config=builder_config,
            podman_client=mock_podman_client,
            audit_logger=mock_audit_logger
        )
        
        # Replace with mocks for tests
        builder.package_validator = Mock()
        builder.uv_installer = Mock()
        builder.image_builder = Mock()
        
        return builder


def test_build_custom_environment_success(builder, tmp_path):
    """Test successful custom environment build."""
    # Mock all sub-components
    builder.package_validator.validate_packages = Mock()
    builder.uv_installer.create_build_context = Mock(return_value=tmp_path)
    
    mock_result = BuildResult(
        success=True,
        image_name="mcp-forge/custom:test-env",
        image_id="sha256:abc123",
        build_time=45.2,
        size_bytes=500_000_000,
        cache_hit=False,
        installed_packages=["numpy==1.24.0"]
    )
    builder.image_builder.build_image = Mock(return_value=mock_result)
    
    result = builder.build_custom_environment(
        name="test-env",
        packages=["numpy>=1.24.0"],
        user_id="user123"
    )
    
    assert result.success is True
    assert result.image_name == "mcp-forge/custom:test-env"
    builder.package_validator.validate_packages.assert_called_once()
    builder.image_builder.build_image.assert_called_once()


def test_build_validates_package_count(builder):
    """Test that build validates package count against maximum."""
    builder.config.max_packages = 10
    
    packages = [f"package{i}" for i in range(20)]
    
    with pytest.raises(ValueError, match="exceeds maximum"):
        builder.build_custom_environment(
            name="test-env",
            packages=packages,
            user_id="user123"
        )


def test_build_validates_packages_with_validator(builder):
    """Test that build uses package validator."""
    builder.package_validator.validate_packages = Mock(
        side_effect=SecurityError("Blocked package")
    )
    
    with pytest.raises(SecurityError, match="Blocked package"):
        builder.build_custom_environment(
            name="test-env",
            packages=["forbidden-package"],
            user_id="user123"
        )


def test_build_enforces_rate_limit(builder):
    """Test that build enforces rate limiting."""
    builder.package_validator.validate_packages = Mock()
    
    # Exhaust rate limit
    for _ in range(5):
        builder.rate_limiter.check_rate_limit("user123")
    
    # Next request should fail
    with pytest.raises(RuntimeError, match="Rate limit exceeded"):
        builder.build_custom_environment(
            name="test-env",
            packages=["numpy"],
            user_id="user123"
        )


def test_build_enforces_concurrent_builds_limit(builder, tmp_path):
    """Test that build enforces concurrent builds limit."""
    builder.config.max_concurrent_builds = 2
    builder.package_validator.validate_packages = Mock()
    builder.uv_installer.create_build_context = Mock(return_value=tmp_path)
    
    # Simulate 2 active builds
    builder.active_builds.add("build1")
    builder.active_builds.add("build2")
    
    with pytest.raises(RuntimeError, match="Maximum concurrent builds"):
        builder.build_custom_environment(
            name="test-env",
            packages=["numpy"],
            user_id="user123"
        )


def test_build_cleans_up_build_context(builder, tmp_path):
    """Test that build cleans up build context even on failure."""
    builder.package_validator.validate_packages = Mock()
    builder.uv_installer.create_build_context = Mock(return_value=tmp_path)
    builder.image_builder.build_image = Mock(side_effect=RuntimeError("Build failed"))
    
    with pytest.raises(RuntimeError, match="Build failed"):
        builder.build_custom_environment(
            name="test-env",
            packages=["numpy"],
            user_id="user123"
        )
    
    # Build context should have been registered for cleanup
    # (actual cleanup would happen in finally block)


def test_build_validates_environment_name(builder):
    """Test that build validates environment name format."""
    with pytest.raises(ValueError, match="must contain only alphanumeric"):
        builder.build_custom_environment(
            name="test env!",  # Invalid: spaces and special chars
            packages=["numpy"],
            user_id="user123"
        )


def test_build_from_template_expands_packages(builder, tmp_path):
    """Test building from template expands package list."""
    builder.package_validator.validate_packages = Mock()
    builder.uv_installer.create_build_context = Mock(return_value=tmp_path)
    
    mock_result = BuildResult(
        success=True,
        image_name="mcp-forge/custom:data-sci",
        image_id="sha256:abc123",
        build_time=45.2,
        size_bytes=500_000_000,
        cache_hit=False,
        installed_packages=["numpy==1.24.0", "pandas==2.0.0", "matplotlib==3.7.0"]
    )
    builder.image_builder.build_image = Mock(return_value=mock_result)
    
    result = builder.build_from_template(
        template_name="data-science",
        name="data-sci",
        user_id="user123"
    )
    
    assert result.success is True
    # Should have called validator with template packages
    builder.package_validator.validate_packages.assert_called_once()
    call_args = builder.package_validator.validate_packages.call_args[0][0]
    assert "numpy" in call_args
    assert "pandas" in call_args
    assert "matplotlib" in call_args


def test_build_from_template_with_additional_packages(builder, tmp_path):
    """Test building from template with additional packages."""
    builder.package_validator.validate_packages = Mock()
    builder.uv_installer.create_build_context = Mock(return_value=tmp_path)
    
    mock_result = BuildResult(
        success=True,
        image_name="mcp-forge/custom:data-sci",
        image_id="sha256:abc123",
        build_time=45.2,
        size_bytes=500_000_000,
        cache_hit=False,
        installed_packages=["numpy==1.24.0", "pandas==2.0.0", "matplotlib==3.7.0", "scipy==1.10.0"]
    )
    builder.image_builder.build_image = Mock(return_value=mock_result)
    
    result = builder.build_from_template(
        template_name="data-science",
        additional_packages=["scipy"],
        name="data-sci",
        user_id="user123"
    )
    
    assert result.success is True
    call_args = builder.package_validator.validate_packages.call_args[0][0]
    assert "scipy" in call_args


def test_build_from_template_validates_template_exists(builder):
    """Test that template build validates template exists."""
    with pytest.raises(ValueError, match="Template.*not found"):
        builder.build_from_template(
            template_name="nonexistent",
            user_id="user123"
        )


def test_list_templates_returns_config_templates(builder):
    """Test that list_templates returns configured templates."""
    templates = builder.list_templates()
    
    assert "data-science" in templates
    assert templates["data-science"]["packages"] == ["numpy", "pandas", "matplotlib"]


def test_rate_limiter_allows_within_limit():
    """Test that rate limiter allows requests within limit."""
    limiter = BuildRateLimiter(max_requests=5, period_seconds=3600)
    
    # Should allow 5 requests
    for _ in range(5):
        limiter.check_rate_limit("user123")
    
    # 6th request should fail
    with pytest.raises(RuntimeError, match="Rate limit exceeded"):
        limiter.check_rate_limit("user123")


def test_rate_limiter_cleans_up_old_requests():
    """Test that rate limiter cleans up old requests."""
    limiter = BuildRateLimiter(max_requests=5, period_seconds=1)
    
    # Make 5 requests
    for _ in range(5):
        limiter.check_rate_limit("user123")
    
    # Wait for period to expire
    import time
    time.sleep(1.1)
    
    # Should allow new request after period
    limiter.check_rate_limit("user123")


def test_rate_limiter_tracks_per_user():
    """Test that rate limiter tracks requests per user."""
    limiter = BuildRateLimiter(max_requests=2, period_seconds=3600)
    
    # User1 makes 2 requests
    limiter.check_rate_limit("user1")
    limiter.check_rate_limit("user1")
    
    # User2 should still be allowed
    limiter.check_rate_limit("user2")
    limiter.check_rate_limit("user2")
    
    # Both users now at limit
    with pytest.raises(RuntimeError):
        limiter.check_rate_limit("user1")
    
    with pytest.raises(RuntimeError):
        limiter.check_rate_limit("user2")


def test_build_registers_and_unregisters_active_builds(builder, tmp_path):
    """Test that builds are registered and unregistered correctly."""
    builder.package_validator.validate_packages = Mock()
    builder.uv_installer.create_build_context = Mock(return_value=tmp_path)
    
    mock_result = BuildResult(
        success=True,
        image_name="mcp-forge/custom:test-env",
        image_id="sha256:abc123",
        build_time=45.2,
        size_bytes=500_000_000,
        cache_hit=False,
        installed_packages=["numpy==1.24.0"]
    )
    builder.image_builder.build_image = Mock(return_value=mock_result)
    
    # Before build
    assert "test-env" not in builder.active_builds
    
    builder.build_custom_environment(
        name="test-env",
        packages=["numpy"],
        user_id="user123"
    )
    
    # After build
    assert "test-env" not in builder.active_builds