379 lines
14 KiB
Python
379 lines
14 KiB
Python
"""
|
|
Tests for Docker Executor component
|
|
|
|
Test suite for Docker-based container execution with isolation,
|
|
resource limits, and audit logging integration.
|
|
"""
|
|
|
|
import pytest
|
|
import tempfile
|
|
from unittest.mock import Mock, patch, MagicMock
|
|
from pathlib import Path
|
|
|
|
# Import components under test
|
|
from src.mai.sandbox.docker_executor import DockerExecutor, ContainerConfig, ContainerResult
|
|
from src.mai.sandbox.audit_logger import AuditLogger
|
|
|
|
|
|
class TestContainerConfig:
|
|
"""Test ContainerConfig dataclass"""
|
|
|
|
def test_default_config(self):
|
|
"""Test default configuration values"""
|
|
config = ContainerConfig()
|
|
assert config.image == "python:3.10-slim"
|
|
assert config.timeout_seconds == 30
|
|
assert config.memory_limit == "128m"
|
|
assert config.cpu_limit == "0.5"
|
|
assert config.network_disabled is True
|
|
assert config.read_only_filesystem is True
|
|
assert config.tmpfs_size == "64m"
|
|
assert config.working_dir == "/app"
|
|
assert config.user == "nobody"
|
|
|
|
def test_custom_config(self):
|
|
"""Test custom configuration values"""
|
|
config = ContainerConfig(
|
|
image="python:3.9-alpine",
|
|
timeout_seconds=60,
|
|
memory_limit="256m",
|
|
cpu_limit="0.8",
|
|
network_disabled=False,
|
|
)
|
|
assert config.image == "python:3.9-alpine"
|
|
assert config.timeout_seconds == 60
|
|
assert config.memory_limit == "256m"
|
|
assert config.cpu_limit == "0.8"
|
|
assert config.network_disabled is False
|
|
|
|
|
|
class TestDockerExecutor:
|
|
"""Test DockerExecutor class"""
|
|
|
|
@pytest.fixture
|
|
def mock_audit_logger(self):
|
|
"""Create mock audit logger"""
|
|
return Mock(spec=AuditLogger)
|
|
|
|
@pytest.fixture
|
|
def docker_executor(self, mock_audit_logger):
|
|
"""Create DockerExecutor instance for testing"""
|
|
return DockerExecutor(audit_logger=mock_audit_logger)
|
|
|
|
def test_init_without_docker(self, mock_audit_logger):
|
|
"""Test initialization when Docker is not available"""
|
|
with patch("src.mai.sandbox.docker_executor.DOCKER_AVAILABLE", False):
|
|
executor = DockerExecutor(audit_logger=mock_audit_logger)
|
|
assert executor.is_available() is False
|
|
assert executor.client is None
|
|
|
|
def test_init_with_docker_error(self, mock_audit_logger):
|
|
"""Test initialization when Docker fails to connect"""
|
|
with patch("src.mai.sandbox.docker_executor.DOCKER_AVAILABLE", True):
|
|
with patch("docker.from_env") as mock_from_env:
|
|
mock_from_env.side_effect = Exception("Docker daemon not running")
|
|
|
|
executor = DockerExecutor(audit_logger=mock_audit_logger)
|
|
assert executor.is_available() is False
|
|
assert executor.client is None
|
|
|
|
def test_is_available(self, docker_executor):
|
|
"""Test is_available method"""
|
|
# When client is None, should not be available
|
|
docker_executor.client = None
|
|
docker_executor.available = False
|
|
assert docker_executor.is_available() is False
|
|
|
|
# When client is available, should reflect available status
|
|
docker_executor.client = Mock()
|
|
docker_executor.available = True
|
|
assert docker_executor.is_available() is True
|
|
|
|
docker_executor.client = Mock()
|
|
docker_executor.available = False
|
|
assert docker_executor.is_available() is False
|
|
|
|
def test_execute_code_unavailable(self, docker_executor):
|
|
"""Test execute_code when Docker is not available"""
|
|
with patch.object(docker_executor, "is_available", return_value=False):
|
|
result = docker_executor.execute_code("print('test')")
|
|
|
|
assert result.success is False
|
|
assert result.container_id == ""
|
|
assert result.exit_code == -1
|
|
assert "Docker executor not available" in result.error
|
|
|
|
@patch("src.mai.sandbox.docker_executor.Path")
|
|
@patch("src.mai.sandbox.docker_executor.tempfile.TemporaryDirectory")
|
|
def test_execute_code_success(self, mock_temp_dir, mock_path, docker_executor):
|
|
"""Test successful code execution in container"""
|
|
# Mock temporary directory and file creation
|
|
mock_temp_file = Mock()
|
|
mock_temp_file.write_text = Mock()
|
|
|
|
mock_temp_path = Mock()
|
|
mock_temp_path.__truediv__ = Mock(return_value=mock_temp_file)
|
|
mock_temp_path.__str__ = Mock(return_value="/tmp/test")
|
|
|
|
mock_temp_dir.return_value.__enter__.return_value = mock_temp_path
|
|
|
|
# Mock Docker client and container
|
|
mock_container = Mock()
|
|
mock_container.id = "test-container-id"
|
|
mock_container.wait.return_value = {"StatusCode": 0}
|
|
mock_container.logs.return_value = b"test output"
|
|
mock_container.stats.return_value = {
|
|
"cpu_stats": {"cpu_usage": {"total_usage": 1000000}, "system_cpu_usage": 2000000},
|
|
"precpu_stats": {"cpu_usage": {"total_usage": 500000}, "system_cpu_usage": 1000000},
|
|
"memory_stats": {"usage": 50000000, "limit": 100000000},
|
|
}
|
|
|
|
mock_client = Mock()
|
|
mock_client.containers.run.return_value = mock_container
|
|
|
|
docker_executor.client = mock_client
|
|
docker_executor.available = True
|
|
|
|
# Execute code
|
|
result = docker_executor.execute_code("print('test')")
|
|
|
|
assert result.success is True
|
|
assert result.container_id == "test-container-id"
|
|
assert result.exit_code == 0
|
|
assert result.stdout == "test output"
|
|
assert result.execution_time > 0
|
|
assert result.resource_usage is not None
|
|
|
|
@patch("src.mai.sandbox.docker_executor.Path")
|
|
@patch("src.mai.sandbox.docker_executor.tempfile.TemporaryDirectory")
|
|
def test_execute_code_with_files(self, mock_temp_dir, mock_path, docker_executor):
|
|
"""Test code execution with additional files"""
|
|
# Mock temporary directory and file creation
|
|
mock_temp_file = Mock()
|
|
mock_temp_file.write_text = Mock()
|
|
|
|
mock_temp_path = Mock()
|
|
mock_temp_path.__truediv__ = Mock(return_value=mock_temp_file)
|
|
mock_temp_path.__str__ = Mock(return_value="/tmp/test")
|
|
|
|
mock_temp_dir.return_value.__enter__.return_value = mock_temp_path
|
|
|
|
# Mock Docker client and container
|
|
mock_container = Mock()
|
|
mock_container.id = "test-container-id"
|
|
mock_container.wait.return_value = {"StatusCode": 0}
|
|
mock_container.logs.return_value = b"test output"
|
|
mock_container.stats.return_value = {}
|
|
|
|
mock_client = Mock()
|
|
mock_client.containers.run.return_value = mock_container
|
|
|
|
docker_executor.client = mock_client
|
|
docker_executor.available = True
|
|
|
|
# Execute code with files
|
|
files = {"data.txt": "test data"}
|
|
result = docker_executor.execute_code("print('test')", files=files)
|
|
|
|
# Verify additional files were handled
|
|
assert mock_temp_file.write_text.call_count >= 2 # code + data file
|
|
assert result.success is True
|
|
|
|
def test_build_container_config(self, docker_executor):
|
|
"""Test building Docker container configuration"""
|
|
config = ContainerConfig(memory_limit="256m", cpu_limit="0.8", network_disabled=False)
|
|
environment = {"TEST_VAR": "test_value"}
|
|
|
|
container_config = docker_executor._build_container_config(config, environment)
|
|
|
|
assert container_config["mem_limit"] == "256m"
|
|
assert container_config["cpu_quota"] == 80000 # 0.8 * 100000
|
|
assert container_config["cpu_period"] == 100000
|
|
assert container_config["network_disabled"] is False
|
|
assert container_config["read_only"] is True
|
|
assert container_config["user"] == "nobody"
|
|
assert container_config["working_dir"] == "/app"
|
|
assert "TEST_VAR" in container_config["environment"]
|
|
assert "security_opt" in container_config
|
|
assert "cap_drop" in container_config
|
|
assert "cap_add" in container_config
|
|
|
|
def test_get_container_stats(self, docker_executor):
|
|
"""Test extracting container resource statistics"""
|
|
mock_container = Mock()
|
|
mock_container.stats.return_value = {
|
|
"cpu_stats": {
|
|
"cpu_usage": {"total_usage": 2000000},
|
|
"system_cpu_usage": 4000000,
|
|
"online_cpus": 2,
|
|
},
|
|
"precpu_stats": {"cpu_usage": {"total_usage": 1000000}, "system_cpu_usage": 2000000},
|
|
"memory_stats": {
|
|
"usage": 67108864, # 64MB
|
|
"limit": 134217728, # 128MB
|
|
},
|
|
}
|
|
|
|
stats = docker_executor._get_container_stats(mock_container)
|
|
|
|
assert stats["cpu_percent"] == 100.0 # (2000000-1000000)/(4000000-2000000) * 2 * 100
|
|
assert stats["memory_usage_bytes"] == 67108864
|
|
assert stats["memory_limit_bytes"] == 134217728
|
|
assert stats["memory_percent"] == 50.0
|
|
assert stats["memory_usage_mb"] == 64.0
|
|
|
|
def test_get_container_stats_error(self, docker_executor):
|
|
"""Test get_container_stats with error"""
|
|
mock_container = Mock()
|
|
mock_container.stats.side_effect = Exception("Stats error")
|
|
|
|
stats = docker_executor._get_container_stats(mock_container)
|
|
|
|
assert stats["cpu_percent"] == 0.0
|
|
assert stats["memory_usage_bytes"] == 0
|
|
assert stats["memory_percent"] == 0.0
|
|
assert stats["memory_usage_mb"] == 0.0
|
|
|
|
def test_log_container_execution(self, docker_executor, mock_audit_logger):
|
|
"""Test logging container execution"""
|
|
config = ContainerConfig(image="python:3.10-slim")
|
|
result = ContainerResult(
|
|
success=True,
|
|
container_id="test-id",
|
|
exit_code=0,
|
|
stdout="test output",
|
|
stderr="",
|
|
execution_time=1.5,
|
|
resource_usage={"cpu_percent": 50.0},
|
|
)
|
|
|
|
docker_executor._log_container_execution("print('test')", result, config)
|
|
|
|
# Verify audit logger was called
|
|
mock_audit_logger.log_execution.assert_called_once()
|
|
call_args = mock_audit_logger.log_execution.call_args
|
|
assert call_args.kwargs["code"] == "print('test')"
|
|
assert call_args.kwargs["execution_type"] == "docker"
|
|
assert "docker_container" in call_args.kwargs["execution_result"]["type"]
|
|
|
|
def test_get_available_images(self, docker_executor):
|
|
"""Test getting available Docker images"""
|
|
mock_image = Mock()
|
|
mock_image.tags = ["python:3.10-slim", "python:3.9-alpine"]
|
|
|
|
mock_client = Mock()
|
|
mock_client.images.list.return_value = [mock_image]
|
|
|
|
docker_executor.client = mock_client
|
|
docker_executor.available = True
|
|
|
|
images = docker_executor.get_available_images()
|
|
|
|
assert "python:3.10-slim" in images
|
|
assert "python:3.9-alpine" in images
|
|
|
|
def test_pull_image(self, docker_executor):
|
|
"""Test pulling Docker image"""
|
|
mock_client = Mock()
|
|
mock_client.images.pull.return_value = None
|
|
|
|
docker_executor.client = mock_client
|
|
docker_executor.available = True
|
|
|
|
result = docker_executor.pull_image("python:3.10-slim")
|
|
|
|
assert result is True
|
|
mock_client.images.pull.assert_called_once_with("python:3.10-slim")
|
|
|
|
def test_cleanup_containers(self, docker_executor):
|
|
"""Test cleaning up containers"""
|
|
mock_container = Mock()
|
|
|
|
mock_client = Mock()
|
|
mock_client.containers.list.return_value = [mock_container, mock_container]
|
|
|
|
docker_executor.client = mock_client
|
|
docker_executor.available = True
|
|
|
|
count = docker_executor.cleanup_containers()
|
|
|
|
assert count == 2
|
|
assert mock_container.remove.call_count == 2
|
|
|
|
def test_get_system_info(self, docker_executor):
|
|
"""Test getting Docker system information"""
|
|
mock_client = Mock()
|
|
mock_client.info.return_value = {
|
|
"Containers": 5,
|
|
"ContainersRunning": 2,
|
|
"Images": 10,
|
|
"MemTotal": 8589934592,
|
|
"NCPU": 4,
|
|
}
|
|
mock_client.version.return_value = {"Version": "20.10.7", "ApiVersion": "1.41"}
|
|
|
|
docker_executor.client = mock_client
|
|
docker_executor.available = True
|
|
|
|
info = docker_executor.get_system_info()
|
|
|
|
assert info["available"] is True
|
|
assert info["version"] == "20.10.7"
|
|
assert info["api_version"] == "1.41"
|
|
assert info["containers"] == 5
|
|
assert info["images"] == 10
|
|
|
|
|
|
class TestDockerExecutorIntegration:
|
|
"""Integration tests for Docker executor with other sandbox components"""
|
|
|
|
@pytest.fixture
|
|
def mock_audit_logger(self):
|
|
"""Create mock audit logger"""
|
|
return Mock(spec=AuditLogger)
|
|
|
|
def test_docker_executor_integration(self, mock_audit_logger):
|
|
"""Test Docker executor integration with audit logger"""
|
|
executor = DockerExecutor(audit_logger=mock_audit_logger)
|
|
|
|
# Test that audit logger is properly integrated
|
|
assert executor.audit_logger is mock_audit_logger
|
|
|
|
# Mock Docker availability for integration test
|
|
with patch.object(executor, "is_available", return_value=False):
|
|
result = executor.execute_code("print('test')")
|
|
|
|
# Should fail gracefully and still attempt logging
|
|
assert result.success is False
|
|
|
|
def test_container_result_serialization(self):
|
|
"""Test ContainerResult can be properly serialized"""
|
|
result = ContainerResult(
|
|
success=True,
|
|
container_id="test-id",
|
|
exit_code=0,
|
|
stdout="test output",
|
|
stderr="",
|
|
execution_time=1.5,
|
|
resource_usage={"cpu_percent": 50.0},
|
|
)
|
|
|
|
# Test that result can be converted to dict for JSON serialization
|
|
result_dict = {
|
|
"success": result.success,
|
|
"container_id": result.container_id,
|
|
"exit_code": result.exit_code,
|
|
"stdout": result.stdout,
|
|
"stderr": result.stderr,
|
|
"execution_time": result.execution_time,
|
|
"error": result.error,
|
|
"resource_usage": result.resource_usage,
|
|
}
|
|
|
|
assert result_dict["success"] is True
|
|
assert result_dict["container_id"] == "test-id"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__])
|