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

SIENTIAPDE-1241: Fixed formatting and linted code

Browse Source
This commit is contained in:
Kou-Kinoshita
2025-10-30 10:38:29 -03:00
parent 17a551d6f9
commit a5c0cb6e47
5 changed files with 161 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
tests/workflows/test_train_model.py
View File

@@ -1,7 +1,8 @@
"""Unit tests for TrainModel workflow."""
from unittest.mock import AsyncMock, Mock, patch
import pytest
from unittest.mock import AsyncMock, MagicMock, Mock, patch
from model_manager.utils.exceptions import ModelTrainingError
from model_manager.utils.models.experiment_status import ExperimentStatus
@@ -116,7 +117,7 @@ def test_extract_error_message_with_cause():
from model_manager.workflows.train_model import TrainModel
workflow_instance = TrainModel()
# Create exception chain
cause = ValueError('Root cause')
exc = RuntimeError('Outer error')
@@ -146,7 +147,7 @@ def test_extract_error_message_circular_reference():
from model_manager.workflows.train_model import TrainModel
workflow_instance = TrainModel()
# Create circular reference
exc1 = ValueError('Error 1')
exc2 = ValueError('Error 2')
@@ -165,7 +166,7 @@ def test_extract_error_message_duplicate_messages():
from model_manager.workflows.train_model import TrainModel
workflow_instance = TrainModel()
# Create chain with duplicate messages
exc1 = ValueError('Same error')
exc2 = ValueError('Same error')
@@ -179,7 +180,9 @@ def test_extract_error_message_duplicate_messages():
@pytest.mark.asyncio
@patch('model_manager.workflows.train_model.workflow')
async def test_validate_training_parameters_success(mock_workflow_module, sample_input_data, mock_train_params):
async def test_validate_training_parameters_success(
mock_workflow_module, sample_input_data, mock_train_params
):
"""Test successful parameter validation."""
from model_manager.workflows.train_model import TrainModel
@@ -189,9 +192,7 @@ async def test_validate_training_parameters_success(mock_workflow_module, sample
workflow_instance = TrainModel()
metadata = {'metadata': {'pod_id': 'test-pod', 'experiment_run_id': 123}}
result = await workflow_instance._validate_training_parameters(
sample_input_data, 123, metadata
)
result = await workflow_instance._validate_training_parameters(sample_input_data, 123, metadata)
assert result == mock_train_params
@@ -214,9 +215,7 @@ async def test_validate_training_parameters_failure(mock_workflow_module, sample
metadata = {'metadata': {'pod_id': 'test-pod', 'experiment_run_id': 123}}
with pytest.raises(ValueError, match='Invalid params'):
await workflow_instance._validate_training_parameters(
sample_input_data, 123, metadata
)
await workflow_instance._validate_training_parameters(sample_input_data, 123, metadata)
# Verify error status update was called
assert mock_workflow_module.execute_activity_method.call_count == 2
@@ -231,7 +230,7 @@ async def test_train_model_success(mock_workflow_module, mock_train_params):
# Setup mocks
train_result = {
'run_name': 'test-run-123',
'run_dir': '/tmp/test-run',
'run_dir': '/tmp/test-run', # noqa: S108
'mse_val': 0.5,
'r2_val': 0.9,
}
@@ -276,10 +275,10 @@ 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')
mock_workflow_module.execute_activity_method = AsyncMock(
side_effect=[mlflow_error, None]
mlflow_error = ModelTrainingError(
model_trained=True, model_saved=False, message='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}}
@@ -308,7 +307,7 @@ async def test_cleanup_resources_success(mock_workflow_module):
await workflow_instance._cleanup_resources(
experiment_run_id=123,
run_dir='/tmp/test-run',
run_dir='/tmp/test-run', # noqa: S108
bucket_name='test-bucket',
file_name='test-file.csv',
metadata=metadata,
@@ -334,7 +333,7 @@ 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',
run_dir='/tmp/test-run', # noqa: S108
bucket_name='test-bucket',
file_name='test-file.csv',
metadata=metadata,
@@ -348,8 +347,8 @@ async def test_cleanup_resources_failure(mock_workflow_module):
@patch('model_manager.workflows.train_model.workflow')
async def test_update_experiment_run_status_only(mock_workflow_module):
"""Test update_experiment_run with status only."""
from model_manager.workflows.train_model import TrainModel
from model_manager.activities.experiment_tracking import UpdateType
from model_manager.workflows.train_model import TrainModel
mock_workflow_module.execute_activity_method = AsyncMock()
@@ -375,8 +374,8 @@ async def test_update_experiment_run_status_only(mock_workflow_module):
@patch('model_manager.workflows.train_model.workflow')
async def test_update_experiment_run_with_error(mock_workflow_module):
"""Test update_experiment_run with error message."""
from model_manager.workflows.train_model import TrainModel
from model_manager.activities.experiment_tracking import UpdateType
from model_manager.workflows.train_model import TrainModel
mock_workflow_module.execute_activity_method = AsyncMock()
@@ -400,8 +399,8 @@ async def test_update_experiment_run_with_error(mock_workflow_module):
@patch('model_manager.workflows.train_model.workflow')
async def test_update_experiment_run_with_run_name(mock_workflow_module):
"""Test update_experiment_run with run_name."""
from model_manager.workflows.train_model import TrainModel
from model_manager.activities.experiment_tracking import UpdateType
from model_manager.workflows.train_model import TrainModel
mock_workflow_module.execute_activity_method = AsyncMock()
@@ -424,24 +423,26 @@ async def test_update_experiment_run_with_run_name(mock_workflow_module):
@pytest.mark.asyncio
@patch('model_manager.workflows.train_model.workflow')
@patch('model_manager.workflows.train_model.POD_ID', 'test-pod-456')
async def test_run_complete_workflow_success(mock_workflow_module, sample_input_data, mock_train_params):
async def test_run_complete_workflow_success(
mock_workflow_module, sample_input_data, mock_train_params
):
"""Test complete workflow execution success path."""
from model_manager.workflows.train_model import TrainModel
# Setup mocks for all activities
train_result = {
'run_name': 'test-run-123',
'run_dir': '/tmp/test-run',
'run_dir': '/tmp/test-run', # noqa: S108
}
mock_workflow_module.execute_activity_method = AsyncMock(
side_effect=[
mock_train_params, # validate_train_params
None, # update status (MAGE_WAITING_PROC)
train_result, # train_model
None, # update status (MLFLOW_SENT)
None, # cleanup_resources
None, # update status (FILE_DELETED)
None, # update status (MAGE_WAITING_PROC)
train_result, # train_model
None, # update status (MLFLOW_SENT)
None, # cleanup_resources
None, # update status (FILE_DELETED)
]
)
@@ -464,7 +465,7 @@ async def test_run_workflow_validation_error(mock_workflow_module, sample_input_
mock_workflow_module.execute_activity_method = AsyncMock(
side_effect=[
ValueError('Invalid parameters'), # validate_train_params fails
None, # update status (ORCHESTRATOR_VALIDATION_ERROR)
None, # update status (ORCHESTRATOR_VALIDATION_ERROR)
]
)
@@ -479,17 +480,19 @@ async def test_run_workflow_validation_error(mock_workflow_module, sample_input_
@pytest.mark.asyncio
@patch('model_manager.workflows.train_model.workflow')
async def test_run_workflow_training_error(mock_workflow_module, sample_input_data, mock_train_params):
async def test_run_workflow_training_error(
mock_workflow_module, sample_input_data, mock_train_params
):
"""Test workflow handles training errors."""
from model_manager.workflows.train_model import TrainModel
# Setup mocks - training fails
mock_workflow_module.execute_activity_method = AsyncMock(
side_effect=[
mock_train_params, # validate_train_params
None, # update status (MAGE_WAITING_PROC)
RuntimeError('Training failed'), # train_model fails
None, # update status (TRAINING_ERROR)
mock_train_params, # validate_train_params
None, # update status (MAGE_WAITING_PROC)
RuntimeError('Training failed'), # train_model fails
None, # update status (TRAINING_ERROR)
]
)
@@ -503,24 +506,26 @@ async def test_run_workflow_training_error(mock_workflow_module, sample_input_da
@pytest.mark.asyncio
@patch('model_manager.workflows.train_model.workflow')
async def test_run_workflow_cleanup_error(mock_workflow_module, sample_input_data, mock_train_params):
async def test_run_workflow_cleanup_error(
mock_workflow_module, sample_input_data, mock_train_params
):
"""Test workflow handles cleanup errors."""
from model_manager.workflows.train_model import TrainModel
# Setup mocks - cleanup fails
train_result = {
'run_name': 'test-run-123',
'run_dir': '/tmp/test-run',
'run_dir': '/tmp/test-run', # noqa: S108
}
mock_workflow_module.execute_activity_method = AsyncMock(
side_effect=[
mock_train_params, # validate_train_params
None, # update status (MAGE_WAITING_PROC)
train_result, # train_model
None, # update status (MLFLOW_SENT)
RuntimeError('Cleanup failed'), # cleanup_resources fails
None, # update status (FILE_DELETE_ERROR)
mock_train_params, # validate_train_params
None, # update status (MAGE_WAITING_PROC)
train_result, # train_model
None, # update status (MLFLOW_SENT)
RuntimeError('Cleanup failed'), # cleanup_resources fails
None, # update status (FILE_DELETE_ERROR)
]
)
@@ -547,13 +552,13 @@ async def test_run_workflow_missing_experiment_run_id():
def test_module_constants():
"""Test that module-level constants are defined correctly."""
from model_manager.workflows.train_model import (
TIMEOUT_VALIDATE_PARAMS,
TIMEOUT_TRAIN_MODEL,
TIMEOUT_DELETE_FILE,
TIMEOUT_TRAIN_MODEL,
TIMEOUT_UPDATE_DATABASE,
TIMEOUT_VALIDATE_PARAMS,
database_retry_policy,
network_retry_policy,
no_retry_policy,
database_retry_policy,
)
# Verify timeouts are integers
@@ -608,11 +613,11 @@ async def test_train_model_empty_run_dir(mock_workflow_module, mock_train_params
mock_workflow_module.execute_activity_method = AsyncMock(
side_effect=[
mock_train_params, # validate
None, # update status
train_result, # train
None, # update status
None, # cleanup
None, # update status
None, # update status
train_result, # train
None, # update status
None, # cleanup
None, # update status
]
)
@@ -644,4 +649,3 @@ async def test_train_model_empty_run_dir(mock_workflow_module, mock_train_params
# Verify cleanup was called with empty string
cleanup_call = mock_workflow_module.execute_activity_method.call_args_list[4]
assert cleanup_call[0][1]['run_dir'] == ''