SIENTIAPDE-1350: Add unit tests for the Cleanup activity, ensuring 100% code coverage, and update todo-list.
This commit is contained in:
554
tests/activities/test_cleanup.py
Normal file
554
tests/activities/test_cleanup.py
Normal file
@@ -0,0 +1,554 @@
|
||||
"""Unit tests for the Cleanup activity, ensuring 100% code coverage."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from importlib import reload
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# Define mocks at the top level to be accessible by all tests
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_logger():
|
||||
"""Fixture for a mock logger."""
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_notification_handler():
|
||||
"""Fixture for a mock notification handler."""
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_metrics_controller():
|
||||
"""Fixture for a mock metrics controller with async methods."""
|
||||
controller = MagicMock()
|
||||
controller.shutdown = AsyncMock()
|
||||
controller.emit = AsyncMock()
|
||||
return controller
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_storage_repository():
|
||||
"""Fixture for a mock storage repository."""
|
||||
repo = MagicMock()
|
||||
repo.delete_file = MagicMock()
|
||||
repo.list_bucket_objects = MagicMock()
|
||||
return repo
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_dir():
|
||||
"""Fixture to create and clean up a temporary directory."""
|
||||
path = tempfile.mkdtemp()
|
||||
yield path
|
||||
shutil.rmtree(path)
|
||||
|
||||
|
||||
# --- Initialization Tests ---
|
||||
|
||||
|
||||
def test_cleanup_init_default_values(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test Cleanup initialization uses default environment values."""
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
assert cleanup.retention_hours == 24
|
||||
assert cleanup.dry_run is False
|
||||
assert cleanup.max_keys_cleanup == 1000
|
||||
|
||||
|
||||
def test_cleanup_init_custom_env_values(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test Cleanup initialization with custom environment values."""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
'CLEANUP_RETENTION_HOURS': '48',
|
||||
'CLEANUP_DRY_RUN': 'true',
|
||||
'MAX_KEYS_CLEANUP': '500',
|
||||
},
|
||||
):
|
||||
import model_manager.activities.cleanup
|
||||
|
||||
reload(model_manager.activities.cleanup)
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
|
||||
assert cleanup.retention_hours == 48
|
||||
assert cleanup.dry_run is True
|
||||
assert cleanup.max_keys_cleanup == 500
|
||||
|
||||
|
||||
@patch.dict(os.environ, {'CLEANUP_RETENTION_HOURS': 'invalid'})
|
||||
def test_cleanup_init_invalid_env_value_raises_error():
|
||||
"""Test Cleanup module raises ValueError for invalid environment variables on import."""
|
||||
import model_manager.activities.cleanup
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
reload(model_manager.activities.cleanup)
|
||||
|
||||
|
||||
# --- MinIO Cleanup Tests ---
|
||||
|
||||
|
||||
def test_cleanup_minio_files_missing_bucket_name(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test cleanup_minio_files raises ValueError if bucket_name is missing."""
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup._emit_metrics = AsyncMock()
|
||||
|
||||
with pytest.raises(ValueError, match='bucket_name must be provided'):
|
||||
asyncio.run(cleanup.cleanup_minio_files({'metadata': {}}))
|
||||
|
||||
|
||||
def test_cleanup_minio_files_success_with_deletions(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test successful deletion of old files from MinIO."""
|
||||
import model_manager.activities.cleanup
|
||||
|
||||
reload(model_manager.activities.cleanup)
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup._emit_metrics = AsyncMock()
|
||||
|
||||
old_ts = int((datetime.now(UTC) - timedelta(hours=48)).timestamp() * 1000)
|
||||
recent_ts = int((datetime.now(UTC) - timedelta(hours=1)).timestamp() * 1000)
|
||||
|
||||
mock_storage_repository.list_bucket_objects.return_value = [
|
||||
f'{old_ts}-old-file.txt',
|
||||
f'{recent_ts}-recent-file.txt',
|
||||
'no-timestamp-file.txt',
|
||||
]
|
||||
|
||||
asyncio.run(cleanup.cleanup_minio_files({'bucket_name': 'test-bucket', 'metadata': {}}))
|
||||
|
||||
mock_storage_repository.delete_file.assert_called_once_with(
|
||||
'test-bucket', f'{old_ts}-old-file.txt'
|
||||
)
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
def test_cleanup_minio_files_dry_run(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test MinIO cleanup in dry_run mode does not delete files."""
|
||||
with patch.dict(os.environ, {'CLEANUP_DRY_RUN': 'true'}):
|
||||
import model_manager.activities.cleanup
|
||||
|
||||
reload(model_manager.activities.cleanup)
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup._emit_metrics = AsyncMock()
|
||||
|
||||
old_ts = int((datetime.now(UTC) - timedelta(hours=48)).timestamp() * 1000)
|
||||
mock_storage_repository.list_bucket_objects.return_value = [f'{old_ts}-old-file.txt']
|
||||
|
||||
asyncio.run(cleanup.cleanup_minio_files({'bucket_name': 'test-bucket', 'metadata': {}}))
|
||||
|
||||
mock_storage_repository.delete_file.assert_not_called()
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
def test_cleanup_minio_files_delete_error(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test error during MinIO file deletion is handled gracefully."""
|
||||
import model_manager.activities.cleanup
|
||||
|
||||
reload(model_manager.activities.cleanup)
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup._emit_metrics = AsyncMock()
|
||||
cleanup.error = AsyncMock()
|
||||
|
||||
old_ts = int((datetime.now(UTC) - timedelta(hours=48)).timestamp() * 1000)
|
||||
mock_storage_repository.list_bucket_objects.return_value = [f'{old_ts}-old-file.txt']
|
||||
mock_storage_repository.delete_file.side_effect = OSError('Permission Denied')
|
||||
|
||||
asyncio.run(cleanup.cleanup_minio_files({'bucket_name': 'test-bucket', 'metadata': {}}))
|
||||
|
||||
cleanup.error.assert_called_once()
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
def test_cleanup_minio_files_exception_handling(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test exception during MinIO cleanup triggers notification and metrics."""
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup.send_notification = AsyncMock()
|
||||
cleanup._emit_metrics = AsyncMock()
|
||||
|
||||
mock_storage_repository.list_bucket_objects.side_effect = Exception('Connection Error')
|
||||
|
||||
with pytest.raises(Exception, match='Connection Error'):
|
||||
asyncio.run(cleanup.cleanup_minio_files({'bucket_name': 'test-bucket', 'metadata': {}}))
|
||||
|
||||
cleanup.send_notification.assert_called_once()
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
# --- Temp Directory Cleanup Tests ---
|
||||
|
||||
|
||||
def test_cleanup_temp_directories_nonexistent_path(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test temp directory cleanup with a non-existent path."""
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup._emit_metrics = AsyncMock()
|
||||
cleanup.warning = AsyncMock()
|
||||
|
||||
asyncio.run(
|
||||
cleanup.cleanup_temp_directories({'temp_path': '/nonexistent/path', 'metadata': {}})
|
||||
)
|
||||
|
||||
cleanup.warning.assert_called_once()
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
def test_cleanup_temp_directories_success_with_deletions(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test successful deletion of old temporary directories."""
|
||||
import model_manager.activities.cleanup
|
||||
|
||||
reload(model_manager.activities.cleanup)
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup._emit_metrics = AsyncMock()
|
||||
|
||||
old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
|
||||
old_dir = os.path.join(temp_dir, f'old_dir_{old_time}')
|
||||
os.makedirs(old_dir)
|
||||
|
||||
recent_time = (datetime.now() - timedelta(hours=1)).strftime('%Y%m%d_%H%M%S_000000')
|
||||
recent_dir = os.path.join(temp_dir, f'recent_dir_{recent_time}')
|
||||
os.makedirs(recent_dir)
|
||||
|
||||
asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))
|
||||
|
||||
assert not os.path.exists(old_dir)
|
||||
assert os.path.exists(recent_dir)
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
@patch.dict(os.environ, {'CLEANUP_DRY_RUN': 'true'})
|
||||
def test_cleanup_temp_directories_dry_run(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test temp directory cleanup in dry_run mode does not delete."""
|
||||
import model_manager.activities.cleanup
|
||||
|
||||
reload(model_manager.activities.cleanup)
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup._emit_metrics = AsyncMock()
|
||||
|
||||
old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
|
||||
old_dir = os.path.join(temp_dir, f'old_dir_{old_time}')
|
||||
os.makedirs(old_dir)
|
||||
|
||||
asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))
|
||||
|
||||
assert os.path.exists(old_dir)
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
def test_cleanup_temp_directories_delete_error(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test graceful handling of errors during directory deletion."""
|
||||
import model_manager.activities.cleanup
|
||||
|
||||
reload(model_manager.activities.cleanup)
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup._emit_metrics = AsyncMock()
|
||||
cleanup.error = AsyncMock()
|
||||
|
||||
old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
|
||||
old_dir = os.path.join(temp_dir, f'old_dir_{old_time}')
|
||||
os.makedirs(old_dir)
|
||||
|
||||
with patch('shutil.rmtree', side_effect=OSError('Permission Denied')):
|
||||
asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))
|
||||
|
||||
cleanup.error.assert_called_once()
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
# --- Metrics and Utility Tests ---
|
||||
|
||||
|
||||
def test_emit_metrics(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test that _emit_metrics calls the public emit_metric method."""
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup.emit_metric = AsyncMock()
|
||||
|
||||
asyncio.run(
|
||||
cleanup._emit_metrics(
|
||||
metadata={'pod_id': 'p1', 'workflow_name': 'wf1'},
|
||||
metrics_status='success',
|
||||
activity_name='test_activity',
|
||||
emit_workflow_metric=True,
|
||||
)
|
||||
)
|
||||
|
||||
assert cleanup.emit_metric.call_count == 2
|
||||
|
||||
|
||||
def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test that files and directories with non-matching names are skipped."""
|
||||
import model_manager.activities.cleanup
|
||||
|
||||
reload(model_manager.activities.cleanup)
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup._emit_metrics = AsyncMock()
|
||||
cleanup.debug = AsyncMock()
|
||||
|
||||
# Create a file and a directory with a non-matching name
|
||||
with open(os.path.join(temp_dir, 'a_file.txt'), 'w') as f:
|
||||
f.write('hello')
|
||||
os.makedirs(os.path.join(temp_dir, 'a_directory_with_no_timestamp'))
|
||||
|
||||
asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))
|
||||
|
||||
# Ensure the debug message for skipping was called for the unmatched directory
|
||||
cleanup.debug.assert_called_with(
|
||||
'Skipping directory without timestamp pattern: a_directory_with_no_timestamp', {}
|
||||
)
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
def test_cleanup_temp_directories_invalid_timestamp_format(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test that a directory with an invalid timestamp format is handled correctly."""
|
||||
import model_manager.activities.cleanup
|
||||
|
||||
reload(model_manager.activities.cleanup)
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup._emit_metrics = AsyncMock()
|
||||
cleanup.error = AsyncMock()
|
||||
|
||||
# Create a directory with a malformed timestamp that matches the regex but fails parsing
|
||||
malformed_dir_name = 'dir_20239999_999999_999999'
|
||||
os.makedirs(os.path.join(temp_dir, malformed_dir_name))
|
||||
|
||||
asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))
|
||||
|
||||
cleanup.error.assert_called_once()
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
def test_cleanup_temp_directories_generic_exception(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test that a generic exception during directory cleanup is handled."""
|
||||
import model_manager.activities.cleanup
|
||||
|
||||
reload(model_manager.activities.cleanup)
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup._emit_metrics = AsyncMock()
|
||||
cleanup.send_notification = AsyncMock()
|
||||
|
||||
with patch('os.listdir', side_effect=Exception('Unexpected OS Error')):
|
||||
with pytest.raises(Exception, match='Unexpected OS Error'):
|
||||
asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))
|
||||
|
||||
cleanup.send_notification.assert_called_once()
|
||||
cleanup._emit_metrics.assert_called_once()
|
||||
|
||||
|
||||
def test_emit_metrics_activity_only(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
):
|
||||
"""Test that _emit_metrics can emit only the activity metric."""
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
|
||||
cleanup = Cleanup(
|
||||
storage_repository=mock_storage_repository,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
metrics_controller=mock_metrics_controller,
|
||||
)
|
||||
cleanup.emit_metric = AsyncMock()
|
||||
|
||||
asyncio.run(
|
||||
cleanup._emit_metrics(
|
||||
metadata={'pod_id': 'p1', 'workflow_name': 'wf1'},
|
||||
metrics_status='success',
|
||||
activity_name='test_activity',
|
||||
emit_workflow_metric=False,
|
||||
)
|
||||
)
|
||||
|
||||
cleanup.emit_metric.assert_called_once()
|
||||
Reference in New Issue
Block a user