gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into release/SIENTIAPDE-1645

Browse Source
This commit is contained in:
vitor-aignosi
2026-04-06 08:34:15 -03:00
parent bf2b6b3888 dfe05a063a
commit d2bb588739
18 changed files with 85 additions and 676 deletions
Download Patch File Download Diff File Expand all files Collapse all files

202
tests/activities/test_cleanup.py
View File

@@ -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,

16
tests/activities/test_training.py
View File

@@ -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('', '')

38
tests/utils/models/test_experiment_status.py
View File

@@ -5,18 +5,15 @@ from model_manager.utils.models.experiment_status import ExperimentStatus
def test_experiment_status_values():
"""Test that all expected status values exist."""
assert ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR == 'ORCHESTRATOR_VALIDATION_ERROR'
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
"""Test that enum has exactly 4 status values."""
assert len(ExperimentStatus) == 4
def test_experiment_status_is_string():
@@ -28,30 +25,25 @@ def test_experiment_status_is_string():
def test_experiment_status_membership():
"""Test membership checks for status values."""
assert 'ORCHESTRATOR_VALIDATION_ERROR' in [s.value for s in ExperimentStatus]
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) == 4
assert ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR in statuses
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
def test_experiment_status_comparison():
"""Test that enum values can be compared with strings."""
assert ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR == 'ORCHESTRATOR_VALIDATION_ERROR'
assert ExperimentStatus.ORCHESTRATOR_WAITING_PROC == 'ORCHESTRATOR_WAITING_PROC'
assert ExperimentStatus.TRAINING_SUCCESS == 'TRAINING_SUCCESS'
assert ExperimentStatus.TRAINING_ERROR != 'TRAINING_SUCCESS'
@@ -59,25 +51,25 @@ def test_experiment_status_comparison():
def test_experiment_status_access_by_name():
"""Test accessing enum members by name."""
assert (
ExperimentStatus['ORCHESTRATOR_VALIDATION_ERROR']
== ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR
)
assert (
ExperimentStatus['ORCHESTRATOR_WAITING_PROC'] == ExperimentStatus.ORCHESTRATOR_WAITING_PROC
)
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
def test_experiment_status_access_by_value():
"""Test accessing enum members by value."""
assert (
ExperimentStatus('ORCHESTRATOR_VALIDATION_ERROR')
== ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR
)
assert (
ExperimentStatus('ORCHESTRATOR_WAITING_PROC') == ExperimentStatus.ORCHESTRATOR_WAITING_PROC
)
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

79
tests/utils/test_exceptions.py
View File

@@ -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) == ''

54
tests/workflows/test_cleanup_files.py
View File

@@ -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',

33
tests/workflows/test_train_model.py
View File

@@ -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,9 +513,8 @@ 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)
None, # update status (TRAINING_SUCCESS)
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