271 lines
8.6 KiB
Python
271 lines
8.6 KiB
Python
"""Tests for Package Validator module."""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock
|
|
|
|
from mcp_forge.builder.package_validator import (
|
|
PackageValidator,
|
|
ApprovalRequiredError,
|
|
SecurityError
|
|
)
|
|
from mcp_forge.config.schema import PackageValidationConfig
|
|
|
|
|
|
@pytest.fixture
|
|
def allowlist_file(tmp_path):
|
|
"""Create temporary allowlist file."""
|
|
allowlist = tmp_path / "allowlist.txt"
|
|
allowlist.write_text("""
|
|
# Standard data science packages
|
|
numpy
|
|
pandas
|
|
scipy
|
|
scikit-learn
|
|
matplotlib
|
|
|
|
# Web and API
|
|
requests
|
|
httpx
|
|
aiohttp
|
|
|
|
# Utilities
|
|
pyyaml
|
|
python-dateutil
|
|
""".strip())
|
|
return allowlist
|
|
|
|
|
|
@pytest.fixture
|
|
def blocklist_file(tmp_path):
|
|
"""Create temporary blocklist file."""
|
|
blocklist = tmp_path / "blocklist.txt"
|
|
blocklist.write_text("""
|
|
# Security concerns
|
|
os-crypto
|
|
subprocess-wrapper
|
|
shell-exec
|
|
|
|
# Known malicious
|
|
malicious-package
|
|
evil-lib
|
|
""".strip())
|
|
return blocklist
|
|
|
|
|
|
@pytest.fixture
|
|
def validation_config(allowlist_file, blocklist_file, tmp_path):
|
|
"""Mock PackageValidationConfig."""
|
|
config = Mock(spec=PackageValidationConfig)
|
|
config.use_allowlist = True
|
|
config.allowlist_path = allowlist_file
|
|
config.blocklist_path = blocklist_file
|
|
config.require_approval_patterns = [
|
|
"^torch.*", # PyTorch packages
|
|
"^tensorflow.*", # TensorFlow packages
|
|
".*-gpu$", # GPU variants
|
|
]
|
|
return config
|
|
|
|
|
|
@pytest.fixture
|
|
def validator(validation_config):
|
|
"""PackageValidator instance."""
|
|
return PackageValidator(validation_config)
|
|
|
|
|
|
def test_extract_package_name_simple(validator):
|
|
"""Test extracting package name from simple spec."""
|
|
assert validator.extract_package_name("numpy") == "numpy"
|
|
assert validator.extract_package_name("pandas") == "pandas"
|
|
|
|
|
|
def test_extract_package_name_with_version(validator):
|
|
"""Test extracting package name with version specifiers."""
|
|
assert validator.extract_package_name("numpy>=1.24.0") == "numpy"
|
|
assert validator.extract_package_name("pandas==2.0.0") == "pandas"
|
|
assert validator.extract_package_name("requests<=2.28.0") == "requests"
|
|
assert validator.extract_package_name("scikit-learn~=1.3.0") == "scikit-learn"
|
|
|
|
|
|
def test_extract_package_name_with_extras(validator):
|
|
"""Test extracting package name with extras."""
|
|
assert validator.extract_package_name("requests[security]") == "requests"
|
|
assert validator.extract_package_name("pandas[excel,sql]") == "pandas"
|
|
|
|
|
|
def test_extract_package_name_complex(validator):
|
|
"""Test extracting package name from complex specs."""
|
|
assert validator.extract_package_name("numpy>=1.24.0,<2.0.0") == "numpy"
|
|
assert validator.extract_package_name("requests[security]>=2.28.0") == "requests"
|
|
|
|
|
|
def test_allowlisted_package_passes(validator):
|
|
"""Test that allowlisted packages pass validation."""
|
|
validator.validate_package("numpy")
|
|
validator.validate_package("pandas>=2.0.0")
|
|
validator.validate_package("requests[security]")
|
|
# Should not raise
|
|
|
|
|
|
def test_blocklisted_package_raises_error(validator):
|
|
"""Test that blocklisted packages raise SecurityError."""
|
|
with pytest.raises(SecurityError, match="malicious-package"):
|
|
validator.validate_package("malicious-package")
|
|
|
|
with pytest.raises(SecurityError, match="evil-lib"):
|
|
validator.validate_package("evil-lib>=1.0.0")
|
|
|
|
|
|
def test_unknown_package_with_allowlist_raises_error(validator):
|
|
"""Test that unknown packages raise error when allowlist is enabled."""
|
|
with pytest.raises(SecurityError, match="unknown-package"):
|
|
validator.validate_package("unknown-package")
|
|
|
|
|
|
def test_package_requiring_approval_raises_error(validator):
|
|
"""Test that packages matching approval patterns raise ApprovalRequiredError."""
|
|
with pytest.raises(ApprovalRequiredError, match="torch"):
|
|
validator.validate_package("torch")
|
|
|
|
with pytest.raises(ApprovalRequiredError, match="tensorflow"):
|
|
validator.validate_package("tensorflow-gpu")
|
|
|
|
with pytest.raises(ApprovalRequiredError, match="gpu"):
|
|
validator.validate_package("cupy-gpu")
|
|
|
|
|
|
def test_validate_packages_list(validator):
|
|
"""Test validating multiple packages at once."""
|
|
packages = ["numpy>=1.24.0", "pandas", "requests"]
|
|
validator.validate_packages(packages)
|
|
# Should not raise
|
|
|
|
|
|
def test_validate_packages_enforces_max_limit(validator):
|
|
"""Test that validate_packages enforces maximum package count."""
|
|
packages = ["numpy", "pandas", "scipy", "matplotlib"]
|
|
|
|
with pytest.raises(ValueError, match="Maximum 3 packages"):
|
|
validator.validate_packages(packages, max_packages=3)
|
|
|
|
|
|
def test_validate_packages_with_mixed_results(validator):
|
|
"""Test that validation stops at first error."""
|
|
packages = ["numpy", "malicious-package", "pandas"]
|
|
|
|
with pytest.raises(SecurityError, match="malicious-package"):
|
|
validator.validate_packages(packages)
|
|
|
|
|
|
def test_validate_packages_with_approval_required(validator):
|
|
"""Test that validation stops at first approval requirement."""
|
|
packages = ["numpy", "torch", "pandas"]
|
|
|
|
with pytest.raises(ApprovalRequiredError, match="torch"):
|
|
validator.validate_packages(packages)
|
|
|
|
|
|
def test_allowlist_loading(allowlist_file):
|
|
"""Test that allowlist is loaded correctly from file."""
|
|
config = Mock(spec=PackageValidationConfig)
|
|
config.use_allowlist = True
|
|
config.allowlist_path = allowlist_file
|
|
config.blocklist_path = None
|
|
config.require_approval_patterns = []
|
|
|
|
validator = PackageValidator(config)
|
|
|
|
assert "numpy" in validator.allowlist
|
|
assert "pandas" in validator.allowlist
|
|
assert "requests" in validator.allowlist
|
|
# Comments and empty lines should be ignored
|
|
assert "# Standard data science packages" not in validator.allowlist
|
|
|
|
|
|
def test_blocklist_loading(blocklist_file):
|
|
"""Test that blocklist is loaded correctly from file."""
|
|
config = Mock(spec=PackageValidationConfig)
|
|
config.use_allowlist = False
|
|
config.allowlist_path = None
|
|
config.blocklist_path = blocklist_file
|
|
config.require_approval_patterns = []
|
|
|
|
validator = PackageValidator(config)
|
|
|
|
assert "malicious-package" in validator.blocklist
|
|
assert "evil-lib" in validator.blocklist
|
|
|
|
|
|
def test_disabled_allowlist_allows_all_except_blocklist(blocklist_file):
|
|
"""Test that disabling allowlist allows any package except blocklisted."""
|
|
config = Mock(spec=PackageValidationConfig)
|
|
config.use_allowlist = False
|
|
config.allowlist_path = None
|
|
config.blocklist_path = blocklist_file
|
|
config.require_approval_patterns = []
|
|
|
|
validator = PackageValidator(config)
|
|
|
|
# Unknown packages should pass
|
|
validator.validate_package("some-random-package")
|
|
|
|
# But blocklisted packages should still fail
|
|
with pytest.raises(SecurityError, match="malicious-package"):
|
|
validator.validate_package("malicious-package")
|
|
|
|
|
|
def test_approval_pattern_matching(validator):
|
|
"""Test that approval patterns match correctly."""
|
|
# torch* should match
|
|
with pytest.raises(ApprovalRequiredError):
|
|
validator.validate_package("torch")
|
|
|
|
with pytest.raises(ApprovalRequiredError):
|
|
validator.validate_package("torchvision")
|
|
|
|
# *-gpu$ should match
|
|
with pytest.raises(ApprovalRequiredError):
|
|
validator.validate_package("something-gpu")
|
|
|
|
|
|
def test_empty_package_list(validator):
|
|
"""Test validating empty package list."""
|
|
validator.validate_packages([])
|
|
# Should not raise
|
|
|
|
|
|
def test_validate_packages_without_max_limit(validator):
|
|
"""Test validating many packages without limit."""
|
|
packages = [f"package{i}" for i in range(100)]
|
|
|
|
# Should raise because packages aren't in allowlist
|
|
with pytest.raises(SecurityError):
|
|
validator.validate_packages(packages)
|
|
|
|
|
|
def test_case_sensitivity(validator):
|
|
"""Test that package names are case-sensitive."""
|
|
# numpy is in allowlist
|
|
validator.validate_package("numpy")
|
|
|
|
# NumPy (different case) should fail
|
|
with pytest.raises(SecurityError, match="NumPy"):
|
|
validator.validate_package("NumPy")
|
|
|
|
|
|
def test_whitespace_handling(validator):
|
|
"""Test that leading/trailing whitespace is handled."""
|
|
validator.validate_package(" numpy ")
|
|
validator.validate_package(" pandas>=2.0.0 ")
|
|
|
|
|
|
def test_file_not_found_handling(tmp_path):
|
|
"""Test handling of missing allowlist/blocklist files."""
|
|
config = Mock(spec=PackageValidationConfig)
|
|
config.use_allowlist = True
|
|
config.allowlist_path = tmp_path / "nonexistent.txt"
|
|
config.blocklist_path = None
|
|
config.require_approval_patterns = []
|
|
|
|
with pytest.raises(FileNotFoundError):
|
|
PackageValidator(config)
|