SIENTIAPDE-1717: Remove MinIO cleanup functionality and associated components. This change streamlines the cleanup workflow to focus solely on local temporary directories, removes the ModelTrainingError exception, and updates related configurations, documentation, and tests.
This commit is contained in:
@@ -4,7 +4,7 @@ import asyncio
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from datetime import datetime, timedelta
|
||||
from importlib import reload
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
@@ -34,15 +34,6 @@ def mock_metrics_controller():
|
||||
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."""
|
||||
@@ -59,11 +50,9 @@ def temp_dir():
|
||||
{
|
||||
'CLEANUP_RETENTION_HOURS': '24',
|
||||
'CLEANUP_DRY_RUN': 'false',
|
||||
'MAX_KEYS_CLEANUP': '1000',
|
||||
},
|
||||
)
|
||||
def test_cleanup_init_default_values(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
@@ -75,7 +64,6 @@ def test_cleanup_init_default_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,
|
||||
@@ -83,11 +71,9 @@ def test_cleanup_init_default_values(
|
||||
|
||||
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,
|
||||
@@ -98,7 +84,6 @@ def test_cleanup_init_custom_env_values(
|
||||
{
|
||||
'CLEANUP_RETENTION_HOURS': '48',
|
||||
'CLEANUP_DRY_RUN': 'true',
|
||||
'MAX_KEYS_CLEANUP': '500',
|
||||
},
|
||||
):
|
||||
import model_manager.activities.cleanup
|
||||
@@ -106,16 +91,14 @@ def test_cleanup_init_custom_env_values(
|
||||
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 = Cleanup(
|
||||
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
|
||||
assert cleanup.retention_hours == 48
|
||||
assert cleanup.dry_run is True
|
||||
|
||||
|
||||
@patch.dict(os.environ, {'CLEANUP_RETENTION_HOURS': 'invalid'})
|
||||
@@ -127,162 +110,10 @@ def test_cleanup_init_invalid_env_value_raises_error():
|
||||
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': {}}))
|
||||
|
||||
|
||||
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
|
||||
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()
|
||||
|
||||
|
||||
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
|
||||
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 = MagicMock()
|
||||
|
||||
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 = MagicMock()
|
||||
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,
|
||||
@@ -291,7 +122,6 @@ def test_cleanup_temp_directories_nonexistent_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,
|
||||
@@ -310,7 +140,6 @@ def test_cleanup_temp_directories_nonexistent_path(
|
||||
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
|
||||
def test_cleanup_temp_directories_success_with_deletions(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
@@ -322,7 +151,6 @@ def test_cleanup_temp_directories_success_with_deletions(
|
||||
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,
|
||||
@@ -347,7 +175,6 @@ def test_cleanup_temp_directories_success_with_deletions(
|
||||
@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,
|
||||
@@ -359,7 +186,6 @@ def test_cleanup_temp_directories_dry_run(
|
||||
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,
|
||||
@@ -379,7 +205,6 @@ def test_cleanup_temp_directories_dry_run(
|
||||
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
|
||||
def test_cleanup_temp_directories_delete_error(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
@@ -391,7 +216,6 @@ def test_cleanup_temp_directories_delete_error(
|
||||
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,
|
||||
@@ -414,7 +238,6 @@ def test_cleanup_temp_directories_delete_error(
|
||||
|
||||
|
||||
def test_emit_metrics(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
@@ -423,7 +246,6 @@ def test_emit_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,
|
||||
@@ -444,7 +266,6 @@ def test_emit_metrics(
|
||||
|
||||
def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
@@ -456,7 +277,6 @@ def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
|
||||
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,
|
||||
@@ -480,7 +300,6 @@ def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
|
||||
|
||||
def test_cleanup_temp_directories_invalid_timestamp_format(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
@@ -492,7 +311,6 @@ def test_cleanup_temp_directories_invalid_timestamp_format(
|
||||
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,
|
||||
@@ -512,7 +330,6 @@ def test_cleanup_temp_directories_invalid_timestamp_format(
|
||||
|
||||
def test_cleanup_temp_directories_generic_exception(
|
||||
temp_dir,
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
@@ -524,7 +341,6 @@ def test_cleanup_temp_directories_generic_exception(
|
||||
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,
|
||||
@@ -541,7 +357,6 @@ def test_cleanup_temp_directories_generic_exception(
|
||||
|
||||
|
||||
def test_emit_metrics_activity_only(
|
||||
mock_storage_repository,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
mock_metrics_controller,
|
||||
@@ -550,7 +365,6 @@ def test_emit_metrics_activity_only(
|
||||
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,
|
||||
|
||||
@@ -342,7 +342,6 @@ def test_train_model_training_fails(
|
||||
):
|
||||
"""Test train_model when training fails."""
|
||||
from model_manager.activities.training import Training
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
@@ -362,11 +361,10 @@ def test_train_model_training_fails(
|
||||
'train_params': mock_train_params,
|
||||
}
|
||||
|
||||
with pytest.raises(ModelTrainingError) as exc_info:
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
asyncio.run(training.train_model(input_data))
|
||||
|
||||
assert exc_info.value.model_trained is False
|
||||
assert exc_info.value.model_saved is False
|
||||
assert str(exc_info.value) == 'Training failed'
|
||||
training.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@@ -382,7 +380,6 @@ def test_train_model_save_fails(
|
||||
):
|
||||
"""Test train_model when model saving fails."""
|
||||
from model_manager.activities.training import Training
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
|
||||
training = Training(
|
||||
model_repository=mock_model_repository,
|
||||
@@ -406,11 +403,10 @@ def test_train_model_save_fails(
|
||||
'train_params': mock_train_params,
|
||||
}
|
||||
|
||||
with pytest.raises(ModelTrainingError) as exc_info:
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
asyncio.run(training.train_model(input_data))
|
||||
|
||||
assert exc_info.value.model_trained is True
|
||||
assert exc_info.value.model_saved is False
|
||||
assert str(exc_info.value) == 'Save failed'
|
||||
training.send_notification.assert_called_once()
|
||||
|
||||
|
||||
@@ -437,14 +433,11 @@ def test_cleanup_resources_success(
|
||||
input_data = {
|
||||
'metadata': {'workflow_id': 'test-123'},
|
||||
'run_dir': '/tmp/run_001', # noqa: S108
|
||||
'bucket_name': 'test-bucket',
|
||||
'file_name': 'test-file.csv',
|
||||
}
|
||||
|
||||
asyncio.run(training.cleanup_resources(input_data))
|
||||
|
||||
mock_model_repository.cleanup_run_directory.assert_called_once_with('/tmp/run_001') # noqa: S108
|
||||
mock_storage_repository.delete_file.assert_called_once_with('test-bucket', 'test-file.csv')
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.TrainingRepository')
|
||||
@@ -510,4 +503,3 @@ def test_cleanup_resources_with_empty_values(
|
||||
asyncio.run(training.cleanup_resources(input_data))
|
||||
|
||||
mock_model_repository.cleanup_run_directory.assert_called_once_with('')
|
||||
mock_storage_repository.delete_file.assert_called_once_with('', '')
|
||||
|
||||
@@ -8,15 +8,12 @@ def test_experiment_status_values():
|
||||
assert ExperimentStatus.ORCHESTRATOR_WAITING_PROC == 'ORCHESTRATOR_WAITING_PROC'
|
||||
assert ExperimentStatus.TRAINING_SUCCESS == 'TRAINING_SUCCESS'
|
||||
assert ExperimentStatus.TRAINING_ERROR == 'TRAINING_ERROR'
|
||||
assert ExperimentStatus.TRACKING_SENT == 'TRACKING_SENT'
|
||||
assert ExperimentStatus.TRACKING_SEND_ERROR == 'TRACKING_SEND_ERROR'
|
||||
assert ExperimentStatus.FILE_DELETED == 'FILE_DELETED'
|
||||
assert ExperimentStatus.FILE_DELETE_ERROR == 'FILE_DELETE_ERROR'
|
||||
|
||||
|
||||
def test_experiment_status_count():
|
||||
"""Test that enum has exactly 8 status values."""
|
||||
assert len(ExperimentStatus) == 8
|
||||
assert len(ExperimentStatus) == 5
|
||||
|
||||
|
||||
def test_experiment_status_is_string():
|
||||
@@ -31,22 +28,16 @@ def test_experiment_status_membership():
|
||||
assert 'ORCHESTRATOR_WAITING_PROC' in [s.value for s in ExperimentStatus]
|
||||
assert 'TRAINING_SUCCESS' in [s.value for s in ExperimentStatus]
|
||||
assert 'TRAINING_ERROR' in [s.value for s in ExperimentStatus]
|
||||
assert 'TRACKING_SENT' in [s.value for s in ExperimentStatus]
|
||||
assert 'TRACKING_SEND_ERROR' in [s.value for s in ExperimentStatus]
|
||||
assert 'FILE_DELETED' in [s.value for s in ExperimentStatus]
|
||||
assert 'FILE_DELETE_ERROR' in [s.value for s in ExperimentStatus]
|
||||
|
||||
|
||||
def test_experiment_status_iteration():
|
||||
"""Test that enum can be iterated."""
|
||||
statuses = list(ExperimentStatus)
|
||||
assert len(statuses) == 8
|
||||
assert len(statuses) == 5
|
||||
assert ExperimentStatus.ORCHESTRATOR_WAITING_PROC in statuses
|
||||
assert ExperimentStatus.TRAINING_SUCCESS in statuses
|
||||
assert ExperimentStatus.TRAINING_ERROR in statuses
|
||||
assert ExperimentStatus.TRACKING_SENT in statuses
|
||||
assert ExperimentStatus.TRACKING_SEND_ERROR in statuses
|
||||
assert ExperimentStatus.FILE_DELETED in statuses
|
||||
assert ExperimentStatus.FILE_DELETE_ERROR in statuses
|
||||
|
||||
|
||||
@@ -64,9 +55,6 @@ def test_experiment_status_access_by_name():
|
||||
)
|
||||
assert ExperimentStatus['TRAINING_SUCCESS'] == ExperimentStatus.TRAINING_SUCCESS
|
||||
assert ExperimentStatus['TRAINING_ERROR'] == ExperimentStatus.TRAINING_ERROR
|
||||
assert ExperimentStatus['TRACKING_SENT'] == ExperimentStatus.TRACKING_SENT
|
||||
assert ExperimentStatus['TRACKING_SEND_ERROR'] == ExperimentStatus.TRACKING_SEND_ERROR
|
||||
assert ExperimentStatus['FILE_DELETED'] == ExperimentStatus.FILE_DELETED
|
||||
assert ExperimentStatus['FILE_DELETE_ERROR'] == ExperimentStatus.FILE_DELETE_ERROR
|
||||
|
||||
|
||||
@@ -77,7 +65,4 @@ def test_experiment_status_access_by_value():
|
||||
)
|
||||
assert ExperimentStatus('TRAINING_SUCCESS') == ExperimentStatus.TRAINING_SUCCESS
|
||||
assert ExperimentStatus('TRAINING_ERROR') == ExperimentStatus.TRAINING_ERROR
|
||||
assert ExperimentStatus('TRACKING_SENT') == ExperimentStatus.TRACKING_SENT
|
||||
assert ExperimentStatus('TRACKING_SEND_ERROR') == ExperimentStatus.TRACKING_SEND_ERROR
|
||||
assert ExperimentStatus('FILE_DELETED') == ExperimentStatus.FILE_DELETED
|
||||
assert ExperimentStatus('FILE_DELETE_ERROR') == ExperimentStatus.FILE_DELETE_ERROR
|
||||
|
||||
@@ -237,83 +237,6 @@ def test_fetch_file_network_error(mock_boto3, mock_logger, storage_config):
|
||||
repo.fetch_file('test-bucket', 'test-file.csv')
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_success(mock_boto3, mock_logger, storage_config):
|
||||
"""Test successful file deletion from MinIO."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.return_value = {}
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Delete file
|
||||
repo.delete_file('test-bucket', 'test-file.csv')
|
||||
|
||||
# Verify delete_object was called correctly
|
||||
mock_s3_client.delete_object.assert_called_once_with(Bucket='test-bucket', Key='test-file.csv')
|
||||
|
||||
# Verify logging
|
||||
assert mock_logger.info.call_count >= 2 # Init + delete
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_non_existent(mock_boto3, mock_logger, storage_config):
|
||||
"""Test delete file that doesn't exist (should succeed silently in S3)."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# S3/MinIO delete is idempotent - deleting non-existent file succeeds
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.return_value = {}
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Delete non-existent file (should succeed)
|
||||
repo.delete_file('test-bucket', 'non-existent.csv')
|
||||
|
||||
mock_s3_client.delete_object.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_access_denied(mock_boto3, mock_logger, storage_config):
|
||||
"""Test delete file when access is denied."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock to raise AccessDenied error
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.side_effect = ClientError(
|
||||
{'Error': {'Code': 'AccessDenied', 'Message': 'Access Denied'}}, 'DeleteObject'
|
||||
)
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Attempt to delete file without permissions
|
||||
with pytest.raises(ClientError) as exc_info:
|
||||
repo.delete_file('test-bucket', 'protected-file.csv')
|
||||
|
||||
assert exc_info.value.response['Error']['Code'] == 'AccessDenied'
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_network_error(mock_boto3, mock_logger, storage_config):
|
||||
"""Test delete file when network error occurs."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock to raise network error
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.side_effect = ConnectionError('Network unreachable')
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Attempt to delete file with network error
|
||||
with pytest.raises(ConnectionError):
|
||||
repo.delete_file('test-bucket', 'test-file.csv')
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_with_ssl(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository initialization with SSL enabled."""
|
||||
@@ -404,24 +327,6 @@ def test_fetch_file_with_special_characters(mock_boto3, mock_logger, storage_con
|
||||
mock_s3_client.get_object.assert_called_once_with(Bucket='test-bucket', Key=special_filename)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_with_path_separators(mock_boto3, mock_logger, storage_config):
|
||||
"""Test delete file with path separators in object key."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.return_value = {}
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Delete file with path separators
|
||||
file_path = 'data/2023/01/test-file.csv'
|
||||
repo.delete_file('test-bucket', file_path)
|
||||
|
||||
mock_s3_client.delete_object.assert_called_once_with(Bucket='test-bucket', Key=file_path)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_different_regions(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository with different AWS regions."""
|
||||
|
||||
@@ -1,79 +0,0 @@
|
||||
"""Unit tests for custom exceptions with 100% coverage."""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_model_training_error_with_default_message():
|
||||
"""Test ModelTrainingError with default message."""
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
|
||||
error = ModelTrainingError(model_trained=True, model_saved=False)
|
||||
|
||||
assert error.model_trained is True
|
||||
assert error.model_saved is False
|
||||
assert str(error) == 'Model training workflow failed (model_trained=True, model_saved=False)'
|
||||
|
||||
|
||||
def test_model_training_error_with_custom_message():
|
||||
"""Test ModelTrainingError with custom message."""
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
|
||||
custom_msg = 'Custom error occurred during training'
|
||||
error = ModelTrainingError(model_trained=False, model_saved=False, message=custom_msg)
|
||||
|
||||
assert error.model_trained is False
|
||||
assert error.model_saved is False
|
||||
assert str(error) == custom_msg
|
||||
|
||||
|
||||
def test_model_training_error_both_true():
|
||||
"""Test ModelTrainingError when both flags are True."""
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
|
||||
error = ModelTrainingError(model_trained=True, model_saved=True)
|
||||
|
||||
assert error.model_trained is True
|
||||
assert error.model_saved is True
|
||||
assert str(error) == 'Model training workflow failed (model_trained=True, model_saved=True)'
|
||||
|
||||
|
||||
def test_model_training_error_both_false():
|
||||
"""Test ModelTrainingError when both flags are False."""
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
|
||||
error = ModelTrainingError(model_trained=False, model_saved=False)
|
||||
|
||||
assert error.model_trained is False
|
||||
assert error.model_saved is False
|
||||
assert str(error) == 'Model training workflow failed (model_trained=False, model_saved=False)'
|
||||
|
||||
|
||||
def test_model_training_error_is_exception():
|
||||
"""Test ModelTrainingError is an Exception subclass."""
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
|
||||
error = ModelTrainingError(model_trained=True, model_saved=False)
|
||||
|
||||
assert isinstance(error, Exception)
|
||||
|
||||
|
||||
def test_model_training_error_can_be_raised():
|
||||
"""Test ModelTrainingError can be raised and caught."""
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
|
||||
with pytest.raises(ModelTrainingError) as exc_info:
|
||||
raise ModelTrainingError(model_trained=True, model_saved=False)
|
||||
|
||||
assert exc_info.value.model_trained is True
|
||||
assert exc_info.value.model_saved is False
|
||||
|
||||
|
||||
def test_model_training_error_with_empty_message():
|
||||
"""Test ModelTrainingError with empty string message."""
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
|
||||
error = ModelTrainingError(model_trained=True, model_saved=True, message='')
|
||||
|
||||
assert error.model_trained is True
|
||||
assert error.model_saved is True
|
||||
assert str(error) == ''
|
||||
@@ -8,8 +8,8 @@ import pytest
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.workflows.cleanup_files.workflow')
|
||||
@patch('model_manager.workflows.cleanup_files.POD_ID', 'temporal-pod')
|
||||
async def test_cleanup_files_workflow_with_input_bucket(mock_workflow_module):
|
||||
"""Test the CleanupFiles workflow when bucket_name is provided in the input."""
|
||||
async def test_cleanup_files_workflow(mock_workflow_module):
|
||||
"""Test the CleanupFiles workflow."""
|
||||
from model_manager.workflows.cleanup_files import CleanupFiles
|
||||
|
||||
# Mock execute_activity_method
|
||||
@@ -17,58 +17,14 @@ async def test_cleanup_files_workflow_with_input_bucket(mock_workflow_module):
|
||||
|
||||
# Instantiate and run the workflow
|
||||
workflow_instance = CleanupFiles()
|
||||
await workflow_instance.run({'bucket_name': 'input-bucket'})
|
||||
await workflow_instance.run({})
|
||||
|
||||
# Verify that the activities were called with the correct parameters
|
||||
calls = mock_workflow_module.execute_activity_method.call_args_list
|
||||
assert len(calls) == 2
|
||||
|
||||
# Check cleanup_minio_files call
|
||||
minio_call_args = calls[0][0][1]
|
||||
assert minio_call_args['bucket_name'] == 'input-bucket'
|
||||
assert minio_call_args['metadata'] == {
|
||||
'pod_id': 'temporal-pod',
|
||||
'workflow_name': 'cleanup_files',
|
||||
}
|
||||
assert len(calls) == 1
|
||||
|
||||
# Check cleanup_temp_directories call
|
||||
local_call_args = calls[1][0][1]
|
||||
assert local_call_args['temp_path'] == 'model_manager/reports/temp'
|
||||
assert local_call_args['metadata'] == {
|
||||
'pod_id': 'temporal-pod',
|
||||
'workflow_name': 'cleanup_files',
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.workflows.cleanup_files.workflow')
|
||||
@patch('model_manager.workflows.cleanup_files.POD_ID', 'temporal-pod')
|
||||
@patch('model_manager.workflows.cleanup_files.DEFAULT_CLEANUP_BUCKET', 'env-var-bucket')
|
||||
async def test_cleanup_files_workflow_with_default_bucket(mock_workflow_module):
|
||||
"""Test the CleanupFiles workflow when using the default bucket from environment variables."""
|
||||
from model_manager.workflows.cleanup_files import CleanupFiles
|
||||
|
||||
# Mock execute_activity_method
|
||||
mock_workflow_module.execute_activity_method = AsyncMock()
|
||||
|
||||
# Instantiate and run the workflow
|
||||
workflow_instance = CleanupFiles()
|
||||
await workflow_instance.run({}) # Empty input
|
||||
|
||||
# Verify that the activities were called
|
||||
calls = mock_workflow_module.execute_activity_method.call_args_list
|
||||
assert len(calls) == 2
|
||||
|
||||
# Check cleanup_minio_files call
|
||||
minio_call_args = calls[0][0][1]
|
||||
assert minio_call_args['bucket_name'] == 'env-var-bucket'
|
||||
assert minio_call_args['metadata'] == {
|
||||
'pod_id': 'temporal-pod',
|
||||
'workflow_name': 'cleanup_files',
|
||||
}
|
||||
|
||||
# Check cleanup_temp_directories call
|
||||
local_call_args = calls[1][0][1]
|
||||
local_call_args = calls[0][0][1]
|
||||
assert local_call_args['temp_path'] == 'model_manager/reports/temp'
|
||||
assert local_call_args['metadata'] == {
|
||||
'pod_id': 'temporal-pod',
|
||||
|
||||
@@ -4,7 +4,6 @@ from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
from model_manager.utils.models.experiment_status import ExperimentStatus
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
|
||||
@@ -275,20 +274,18 @@ async def test_train_model_mlflow_error(mock_workflow_module, mock_train_params)
|
||||
from model_manager.workflows.train_model import TrainModel
|
||||
|
||||
# Setup mocks - MLflow save fails
|
||||
mlflow_error = ModelTrainingError(
|
||||
model_trained=True, model_saved=False, message='MLflow save failed'
|
||||
)
|
||||
mlflow_error = RuntimeError('MLflow save failed')
|
||||
mock_workflow_module.execute_activity_method = AsyncMock(side_effect=[mlflow_error, None])
|
||||
|
||||
workflow_instance = TrainModel()
|
||||
metadata = {'metadata': {'pod_id': 'test-pod', 'experiment_run_id': 123}}
|
||||
|
||||
with pytest.raises(ModelTrainingError):
|
||||
with pytest.raises(RuntimeError):
|
||||
await workflow_instance._train_model(mock_train_params, 123, metadata)
|
||||
|
||||
# Verify TRACKING_SEND_ERROR status was set
|
||||
# Verify TRAINING_ERROR status was set
|
||||
call_args = mock_workflow_module.execute_activity_method.call_args_list[1]
|
||||
assert call_args[0][1]['status'] == ExperimentStatus.TRACKING_SEND_ERROR
|
||||
assert call_args[0][1]['status'] == ExperimentStatus.TRAINING_ERROR
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -306,14 +303,11 @@ async def test_cleanup_resources_success(mock_workflow_module):
|
||||
metadata = {'metadata': {'pod_id': 'test-pod', 'experiment_run_id': 123}}
|
||||
|
||||
await workflow_instance._cleanup_resources(
|
||||
experiment_run_id=123,
|
||||
run_dir='/tmp/test-run', # noqa: S108
|
||||
bucket_name='test-bucket',
|
||||
file_name='test-file.csv',
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
assert mock_workflow_module.execute_activity_method.call_count == 2
|
||||
assert mock_workflow_module.execute_activity_method.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -332,15 +326,12 @@ async def test_cleanup_resources_failure(mock_workflow_module):
|
||||
|
||||
with pytest.raises(RuntimeError, match='Cleanup failed'):
|
||||
await workflow_instance._cleanup_resources(
|
||||
experiment_run_id=123,
|
||||
run_dir='/tmp/test-run', # noqa: S108
|
||||
bucket_name='test-bucket',
|
||||
file_name='test-file.csv',
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
# Verify error status update was called
|
||||
assert mock_workflow_module.execute_activity_method.call_count == 2
|
||||
assert mock_workflow_module.execute_activity_method.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -411,7 +402,7 @@ async def test_update_experiment_run_with_run_name(mock_workflow_module):
|
||||
metadata=metadata,
|
||||
experiment_run_id=123,
|
||||
update_type=UpdateType.MODEL_SAVED,
|
||||
status=ExperimentStatus.TRACKING_SENT,
|
||||
status=ExperimentStatus.TRAINING_SUCCESS,
|
||||
run_name='test-run-123',
|
||||
)
|
||||
|
||||
@@ -440,9 +431,8 @@ async def test_run_complete_workflow_success(
|
||||
mock_train_params, # validate_train_params
|
||||
None, # update status (ORCHESTRATOR_WAITING_PROC)
|
||||
train_result, # train_model
|
||||
None, # update status (TRACKING_SENT)
|
||||
None, # update status (TRAINING_SUCCESS)
|
||||
None, # cleanup_resources
|
||||
None, # update status (FILE_DELETED)
|
||||
]
|
||||
)
|
||||
|
||||
@@ -452,7 +442,7 @@ async def test_run_complete_workflow_success(
|
||||
await workflow_instance.run(sample_input_data)
|
||||
|
||||
# Verify all activities were called
|
||||
assert mock_workflow_module.execute_activity_method.call_count == 6
|
||||
assert mock_workflow_module.execute_activity_method.call_count == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -523,7 +513,6 @@ async def test_run_workflow_cleanup_error(
|
||||
mock_train_params, # validate_train_params
|
||||
None, # update status (ORCHESTRATOR_WAITING_PROC)
|
||||
train_result, # train_model
|
||||
None, # update status (TRACKING_SENT)
|
||||
RuntimeError('Cleanup failed'), # cleanup_resources fails
|
||||
None, # update status (FILE_DELETE_ERROR)
|
||||
]
|
||||
@@ -534,7 +523,7 @@ async def test_run_workflow_cleanup_error(
|
||||
with pytest.raises(RuntimeError, match='Cleanup failed'):
|
||||
await workflow_instance.run(sample_input_data)
|
||||
|
||||
assert mock_workflow_module.execute_activity_method.call_count == 6
|
||||
assert mock_workflow_module.execute_activity_method.call_count == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
Reference in New Issue
Block a user