"""Tests for the resource cleanup manager.""" import pytest import tempfile import time from pathlib import Path from unittest.mock import Mock, patch, MagicMock from datetime import datetime, timezone, timedelta from src.services.resource_cleanup_manager import ResourceCleanupManager, CleanupStats class TestResourceCleanupManager: """Test cases for ResourceCleanupManager.""" @pytest.fixture def cleanup_manager(self): """Create a ResourceCleanupManager instance for testing.""" return ResourceCleanupManager(cleanup_interval_seconds=1, max_temp_files=10) @pytest.fixture def temp_dir(self): """Create a temporary directory for testing.""" with tempfile.TemporaryDirectory() as temp_dir: yield Path(temp_dir) def test_initialization(self, cleanup_manager): """Test cleanup manager initialization.""" assert cleanup_manager.cleanup_interval_seconds == 1 assert cleanup_manager.max_temp_files == 10 assert isinstance(cleanup_manager.temp_directories, set) assert isinstance(cleanup_manager.cache_directories, set) assert isinstance(cleanup_manager.model_references, set) assert cleanup_manager._cleanup_thread.is_alive() def test_register_temp_directory(self, cleanup_manager, temp_dir): """Test registering temporary directory.""" cleanup_manager.register_temp_directory(temp_dir) assert temp_dir in cleanup_manager.temp_directories def test_register_cache_directory(self, cleanup_manager, temp_dir): """Test registering cache directory.""" cleanup_manager.register_cache_directory(temp_dir) assert temp_dir in cleanup_manager.cache_directories def test_register_model_reference(self, cleanup_manager): """Test registering model reference.""" mock_model = Mock() cleanup_manager.register_model_reference(mock_model) assert mock_model in cleanup_manager.model_references @patch('src.services.resource_cleanup_manager.psutil.virtual_memory') @patch('src.services.resource_cleanup_manager.torch.cuda.is_available') @patch('src.services.resource_cleanup_manager.torch.cuda.memory_allocated') @patch('src.services.resource_cleanup_manager.torch.cuda.empty_cache') @patch('src.services.resource_cleanup_manager.gc.collect') def test_cleanup_memory(self, mock_gc_collect, mock_empty_cache, mock_memory_allocated, mock_cuda_available, mock_virtual_memory, cleanup_manager): """Test memory cleanup.""" # Mock memory usage mock_memory = Mock() mock_memory.used = 8 * 1024 * 1024 # 8MB initial mock_virtual_memory.return_value = mock_memory mock_cuda_available.return_value = True mock_memory_allocated.return_value = 4 * 1024 * 1024 # 4MB GPU memory mock_gc_collect.return_value = 10 # 10 objects collected # Add some model references mock_model1 = Mock() mock_model2 = Mock() cleanup_manager.register_model_reference(mock_model1) cleanup_manager.register_model_reference(mock_model2) # Perform cleanup result = cleanup_manager.cleanup_memory() # Verify cleanup was performed assert "memory_freed_mb" in result assert "gpu_memory_freed_mb" in result assert "objects_collected" in result assert "cleanup_duration_seconds" in result # Verify model references were cleared assert len(cleanup_manager.model_references) == 0 # Verify garbage collection was called mock_gc_collect.assert_called_once() # Verify GPU cache was cleared mock_empty_cache.assert_called_once() def test_cleanup_temp_files(self, cleanup_manager, temp_dir): """Test temporary file cleanup.""" # Create some test files test_file1 = temp_dir / "test1.txt" test_file2 = temp_dir / "test2.txt" test_file3 = temp_dir / "test3.txt" test_file1.write_text("test content 1") test_file2.write_text("test content 2") test_file3.write_text("test content 3") # Set file modification times old_time = datetime.now(timezone.utc) - timedelta(hours=2) # 2 hours ago new_time = datetime.now(timezone.utc) - timedelta(minutes=30) # 30 minutes ago # Make two files old test_file1.touch() test_file2.touch() test_file3.touch() # Register temp directory cleanup_manager.register_temp_directory(temp_dir) # Perform cleanup result = cleanup_manager.cleanup_temp_files() # Verify cleanup statistics assert "files_removed" in result assert "directories_removed" in result assert "cleanup_duration_seconds" in result # Note: In real scenarios, files would be removed based on modification time # This test verifies the cleanup logic structure def test_cleanup_cache(self, cleanup_manager, temp_dir): """Test cache cleanup.""" # Create some test cache files cache_file1 = temp_dir / "cache1.pkl" cache_file2 = temp_dir / "cache2.pkl" cache_file1.write_text("cache content 1") cache_file2.write_text("cache content 2") # Register cache directory cleanup_manager.register_cache_directory(temp_dir) # Perform cleanup result = cleanup_manager.cleanup_cache() # Verify cleanup statistics assert "cache_entries_cleared" in result assert "cleanup_duration_seconds" in result @patch('src.services.resource_cleanup_manager.psutil.virtual_memory') @patch('src.services.resource_cleanup_manager.torch.cuda.is_available') @patch('src.services.resource_cleanup_manager.torch.cuda.memory_allocated') @patch('src.services.resource_cleanup_manager.torch.cuda.empty_cache') @patch('src.services.resource_cleanup_manager.gc.collect') def test_perform_full_cleanup(self, mock_gc_collect, mock_empty_cache, mock_memory_allocated, mock_cuda_available, mock_virtual_memory, cleanup_manager, temp_dir): """Test full cleanup operation.""" # Mock memory usage mock_memory = Mock() mock_memory.used = 8 * 1024 * 1024 # 8MB mock_virtual_memory.return_value = mock_memory mock_cuda_available.return_value = True mock_memory_allocated.return_value = 4 * 1024 * 1024 # 4MB GPU memory mock_gc_collect.return_value = 5 # 5 objects collected # Register directories and models cleanup_manager.register_temp_directory(temp_dir) cleanup_manager.register_cache_directory(temp_dir) cleanup_manager.register_model_reference(Mock()) # Perform full cleanup stats = cleanup_manager.perform_full_cleanup() # Verify cleanup stats assert isinstance(stats, CleanupStats) assert stats.memory_freed_mb >= 0 assert stats.gpu_memory_freed_mb >= 0 assert stats.temp_files_removed >= 0 assert stats.cache_entries_cleared >= 0 assert stats.cleanup_duration_seconds >= 0 assert isinstance(stats.timestamp, datetime) # Verify stats were added to history assert len(cleanup_manager.cleanup_stats) > 0 assert cleanup_manager.cleanup_stats[-1] == stats def test_get_cleanup_stats(self, cleanup_manager): """Test getting cleanup statistics.""" # Add some mock stats mock_stats = [ CleanupStats( memory_freed_mb=100.0, gpu_memory_freed_mb=50.0, temp_files_removed=5, cache_entries_cleared=10, cleanup_duration_seconds=1.5, timestamp=datetime.now(timezone.utc) ), CleanupStats( memory_freed_mb=200.0, gpu_memory_freed_mb=75.0, temp_files_removed=8, cache_entries_cleared=15, cleanup_duration_seconds=2.0, timestamp=datetime.now(timezone.utc) ) ] cleanup_manager.cleanup_stats = mock_stats # Get stats with limit stats = cleanup_manager.get_cleanup_stats(limit=1) assert len(stats) == 1 assert stats[0] == mock_stats[-1] @patch('src.services.resource_cleanup_manager.psutil.virtual_memory') @patch('src.services.resource_cleanup_manager.torch.cuda.is_available') @patch('src.services.resource_cleanup_manager.torch.cuda.memory_allocated') @patch('src.services.resource_cleanup_manager.torch.cuda.get_device_properties') def test_get_memory_usage(self, mock_gpu_props, mock_memory_allocated, mock_cuda_available, mock_virtual_memory, cleanup_manager): """Test memory usage monitoring.""" # Mock system memory mock_memory = Mock() mock_memory.total = 16 * 1024**3 # 16GB mock_memory.available = 8 * 1024**3 # 8GB mock_memory.used = 8 * 1024**3 # 8GB mock_memory.percent = 50.0 mock_virtual_memory.return_value = mock_memory # Mock GPU memory mock_cuda_available.return_value = True mock_memory_allocated.return_value = 4 * 1024**3 # 4GB allocated mock_gpu_props.return_value.total_memory = 8 * 1024**3 # 8GB total # Get memory usage usage = cleanup_manager.get_memory_usage() # Verify system memory assert usage["total_memory_gb"] == 16.0 assert usage["available_memory_gb"] == 8.0 assert usage["used_memory_gb"] == 8.0 assert usage["memory_percent"] == 50.0 # Verify GPU memory assert usage["gpu_memory_allocated_gb"] == 4.0 assert usage["gpu_memory_total_gb"] == 8.0 assert usage["gpu_memory_percent"] == 50.0 def test_should_perform_cleanup_high_memory(self, cleanup_manager): """Test cleanup decision with high memory usage.""" # Mock high memory usage with patch.object(cleanup_manager, 'get_memory_usage') as mock_usage: mock_usage.return_value = { "memory_percent": 85.0, # High memory usage "gpu_memory_percent": 30.0 } should_cleanup = cleanup_manager.should_perform_cleanup() assert should_cleanup is True def test_should_perform_cleanup_high_gpu_memory(self, cleanup_manager): """Test cleanup decision with high GPU memory usage.""" # Mock high GPU memory usage with patch.object(cleanup_manager, 'get_memory_usage') as mock_usage: mock_usage.return_value = { "memory_percent": 50.0, "gpu_memory_percent": 85.0 # High GPU memory usage } should_cleanup = cleanup_manager.should_perform_cleanup() assert should_cleanup is True def test_should_perform_cleanup_time_based(self, cleanup_manager): """Test cleanup decision based on time since last cleanup.""" # Set last cleanup to be a long time ago cleanup_manager.last_cleanup = datetime.now(timezone.utc) - timedelta(hours=2) # Mock normal memory usage with patch.object(cleanup_manager, 'get_memory_usage') as mock_usage: mock_usage.return_value = { "memory_percent": 50.0, "gpu_memory_percent": 30.0 } should_cleanup = cleanup_manager.should_perform_cleanup() assert should_cleanup is True def test_should_perform_cleanup_normal_usage(self, cleanup_manager): """Test cleanup decision with normal memory usage.""" # Mock normal memory usage with patch.object(cleanup_manager, 'get_memory_usage') as mock_usage: mock_usage.return_value = { "memory_percent": 50.0, "gpu_memory_percent": 30.0 } should_cleanup = cleanup_manager.should_perform_cleanup() # Should not cleanup unless it's been a long time assert should_cleanup is False def test_shutdown(self, cleanup_manager): """Test cleanup manager shutdown.""" # Verify thread is running assert cleanup_manager._cleanup_thread.is_alive() # Shutdown cleanup_manager.shutdown() # Verify thread is stopped assert not cleanup_manager._cleanup_thread.is_alive() def test_cleanup_stats_dataclass(self): """Test CleanupStats dataclass.""" timestamp = datetime.now(timezone.utc) stats = CleanupStats( memory_freed_mb=100.0, gpu_memory_freed_mb=50.0, temp_files_removed=5, cache_entries_cleared=10, cleanup_duration_seconds=1.5, timestamp=timestamp ) assert stats.memory_freed_mb == 100.0 assert stats.gpu_memory_freed_mb == 50.0 assert stats.temp_files_removed == 5 assert stats.cache_entries_cleared == 10 assert stats.cleanup_duration_seconds == 1.5 assert stats.timestamp == timestamp