"""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 mcp_forge.podman.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