test_complete_system.py

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!")