import unittest import asyncio import tempfile import shutil from pathlib import Path from datetime import datetime from data_models import SpaceInfo, ErrorInfo, RepairStrategy, SpaceStatus, ErrorType, RepairAction from safety_validator import SafetyValidator, RiskLevel, ValidationResult from rollback_manager import RollbackManager, BackupType class TestSafetyValidator(unittest.TestCase): def setUp(self): self.validator = SafetyValidator() self.temp_dir = Path(tempfile.mkdtemp()) def tearDown(self): shutil.rmtree(self.temp_dir, ignore_errors=True) def test_security_scan(self): async def run_test(): malicious_code = ''' import os eval(user_input) os.system("rm -rf /") api_key = "sk-1234567890abcdef" ''' test_file = self.temp_dir / "malicious.py" test_file.write_text(malicious_code) issues = await self.validator.security_checker.scan_code(str(test_file), malicious_code) self.assertGreater(len(issues), 0) issue_types = [issue.issue_type for issue in issues] self.assertIn('command_injection', issue_types) self.assertIn('hardcoded_secrets', issue_types) asyncio.run(run_test()) def test_risk_assessment(self): async def run_test(): space_info = SpaceInfo( space_id="test/space", name="test", repository_url="https://huggingface.co/spaces/test/space", current_status=SpaceStatus.ERROR, last_updated=datetime.now() ) strategy = RepairStrategy( action=RepairAction.UPDATE_DEPENDENCIES, description="Update dependencies", modifications={"type": "test"}, risk_level="medium" ) result = await self.validator.risk_assessor.assess_risk( space_info, strategy, ["requirements.txt"] ) self.assertIsInstance(result, ValidationResult) self.assertIn(result.status, ['passed', 'warning', 'failed']) self.assertIsInstance(result.risk_level, RiskLevel) self.assertGreaterEqual(result.confidence, 0.0) self.assertLessEqual(result.confidence, 1.0) asyncio.run(run_test()) class TestRollbackManager(unittest.TestCase): def setUp(self): self.temp_dir = Path(tempfile.mkdtemp()) self.rollback_manager = RollbackManager(str(self.temp_dir / "backups")) self.test_file = self.temp_dir / "test.txt" self.test_content = "original content" self.test_file.write_text(self.test_content) def tearDown(self): shutil.rmtree(self.temp_dir, ignore_errors=True) def test_backup_and_restore(self): async def run_test(): backup_id = await self.rollback_manager.backup_strategy.create_backup( "test/space", str(self.test_file), BackupType.FILE, "Test backup" ) self.assertIsNotNone(backup_id) self.assertTrue(backup_id.startswith("test/space_file_")) backup_info = await self.rollback_manager.backup_strategy.get_backup_info(backup_id) self.assertIsNotNone(backup_info) self.assertEqual(backup_info.space_id, "test/space") self.assertEqual(backup_info.backup_type, BackupType.FILE) self.test_file.write_text("modified content") success = await self.rollback_manager.execute_rollback(backup_id) self.assertTrue(success) restored_content = self.test_file.read_text() self.assertEqual(restored_content, self.test_content) asyncio.run(run_test()) class TestRepairStrategy(unittest.TestCase): def test_repair_strategy_creation(self): strategy = RepairStrategy( action=RepairAction.MODIFY_DOCKERFILE, description="Fix Dockerfile syntax", modifications={ "type": "syntax_fix", "target_line": 3, "new_line": "RUN pip install --no-cache-dir -r requirements.txt" }, risk_level="low", success_rate=0.9, estimated_time=180, rollback_possible=True ) self.assertEqual(strategy.action, RepairAction.MODIFY_DOCKERFILE) self.assertEqual(strategy.description, "Fix Dockerfile syntax") self.assertEqual(strategy.risk_level, "low") self.assertEqual(strategy.success_rate, 0.9) self.assertEqual(strategy.estimated_time, 180) self.assertTrue(strategy.rollback_possible) self.assertFalse(strategy.manual_review_required) class TestErrorInfo(unittest.TestCase): def test_error_info_creation(self): error_info = ErrorInfo( error_type=ErrorType.DEPENDENCY_INSTALL, message="pip install failed", log_snippet="ERROR: Could not find a version", confidence=0.95, occurred_at=datetime.now(), suggested_fixes=["update requirements.txt", "change pip source"] ) self.assertEqual(error_info.error_type, ErrorType.DEPENDENCY_INSTALL) self.assertEqual(error_info.message, "pip install failed") self.assertEqual(error_info.confidence, 0.95) self.assertEqual(len(error_info.suggested_fixes), 2) self.assertIsInstance(error_info.occurred_at, datetime) class TestSpaceInfo(unittest.TestCase): def test_space_info_creation(self): space_info = SpaceInfo( space_id="test/demo-space", name="demo-space", repository_url="https://huggingface.co/spaces/test/demo-space", description="A demo space for testing", author="test-user", tags=["demo", "test"], sdk="gradio", python_version="3.9", dockerfile_path="Dockerfile", local_path="/tmp/demo-space", created_at=datetime.now(), last_modified=datetime.now() ) self.assertEqual(space_info.space_id, "test/demo-space") self.assertEqual(space_info.name, "demo-space") self.assertEqual(space_info.sdk, "gradio") self.assertEqual(space_info.python_version, "3.9") self.assertEqual(len(space_info.tags), 2) self.assertIn("demo", space_info.tags) self.assertIn("test", space_info.tags) def run_integration_tests(): async def integration_test(): temp_dir = Path(tempfile.mkdtemp()) try: validator = SafetyValidator() rollback_manager = RollbackManager(str(temp_dir / "backups")) space_info = SpaceInfo( space_id="integration/test", name="integration-test", repository_url="https://huggingface.co/spaces/integration/test", current_status=SpaceStatus.ERROR, last_updated=datetime.now() ) error_info = ErrorInfo( error_type=ErrorType.DEPENDENCY_INSTALL, message="Integration test error", confidence=0.8 ) strategy = RepairStrategy( action=RepairAction.UPDATE_DEPENDENCIES, description="Integration test strategy", modifications={"type": "dependency_update"}, risk_level="low" ) safety_result = await validator.validate_repair_safety( space_info, error_info, strategy, ["requirements.txt"] ) assert isinstance(safety_result, ValidationResult) assert safety_result.status.value in ['passed', 'warning', 'failed'] test_file = temp_dir / "requirements.txt" test_file.write_text("torch==1.9.0\ntransformers>=4.0.0") backup_id = await rollback_manager.backup_strategy.create_backup( "integration/test", str(test_file), BackupType.FILE, "Integration test backup" ) assert backup_id is not None test_file.write_text("torch==2.0.0\ntransformers>=4.20.0") rollback_success = await rollback_manager.execute_rollback(backup_id) assert rollback_success restored_content = test_file.read_text() assert "torch==1.9.0" in restored_content print("✅ Integration tests passed") finally: shutil.rmtree(temp_dir, ignore_errors=True) asyncio.run(integration_test()) if __name__ == "__main__": print("🧪 Running unit tests...") unittest.main(argv=[''], exit=False, verbosity=2) print("\n🔗 Running integration tests...") run_integration_tests() print("\n✨ All tests completed!")