355 lines
11 KiB
Python
355 lines
11 KiB
Python
"""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=[]
|
|
)
|