337 lines
13 KiB
Python
337 lines
13 KiB
Python
"""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
|