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

16
tests/test_metrics.py
View File

@@ -4,10 +4,6 @@ This module tests the Prometheus metrics configuration used for
monitoring and observability in the Sientia DataOps Model Manager. monitoring and observability in the Sientia DataOps Model Manager.
""" """
from unittest.mock import MagicMock, patch
import pytest
def test_app_up_metric_exists(): def test_app_up_metric_exists():
"""Test that APP_UP metric is properly defined.""" """Test that APP_UP metric is properly defined."""
@@ -15,7 +11,9 @@ def test_app_up_metric_exists():
assert APP_UP is not None assert APP_UP is not None
assert APP_UP._name == 'app_up' assert APP_UP._name == 'app_up'
assert APP_UP._documentation == 'Indicates if the application is running (1) or shutting down (0)' assert (
APP_UP._documentation == 'Indicates if the application is running (1) or shutting down (0)'
)
def test_app_up_metric_has_pod_id_label(): def test_app_up_metric_has_pod_id_label():
@@ -75,11 +73,11 @@ def test_app_up_metric_multiple_pods():
def test_app_up_metric_default_value(): def test_app_up_metric_default_value():
"""Test that APP_UP metric starts with no value set.""" """Test that APP_UP metric starts with no value set."""
from model_manager.metrics import APP_UP
# Create a new label that hasn't been used yet # Create a new label that hasn't been used yet
import uuid import uuid
from model_manager.metrics import APP_UP
unique_pod = f'test-pod-{uuid.uuid4()}' unique_pod = f'test-pod-{uuid.uuid4()}'
# The metric should exist but not have a value until set # The metric should exist but not have a value until set
@@ -204,8 +202,8 @@ def test_app_up_metric_thread_safety():
def test_prometheus_client_gauge_import(): def test_prometheus_client_gauge_import():
"""Test that Gauge is properly imported from prometheus_client.""" """Test that Gauge is properly imported from prometheus_client."""
from model_manager.metrics import Gauge
from prometheus_client import Gauge as PrometheusGauge from prometheus_client import Gauge as PrometheusGauge
assert Gauge is PrometheusGauge from model_manager.metrics import Gauge
assert Gauge is PrometheusGauge

3
tests/utils/repository/test_storage_repository.py
View File

@@ -1,7 +1,7 @@
"""Unit tests for StorageRepository class.""" """Unit tests for StorageRepository class."""
from io import BytesIO from io import BytesIO
from unittest.mock import MagicMock, Mock, patch from unittest.mock import Mock, patch
import pytest import pytest
from botocore.exceptions import ClientError from botocore.exceptions import ClientError
@@ -462,4 +462,3 @@ def test_fetch_file_logs_file_size(mock_boto3, mock_logger, storage_config):
# Verify logging includes file size # Verify logging includes file size
log_calls = [str(call) for call in mock_logger.info.call_args_list] log_calls = [str(call) for call in mock_logger.info.call_args_list]
assert any('12345 bytes' in str(call) for call in log_calls) assert any('12345 bytes' in str(call) for call in log_calls)

116
tests/utils/repository/test_training_repository.py
View File

@@ -496,12 +496,14 @@ class TestTrain:
): ):
"""Test basic training workflow.""" """Test basic training workflow."""
# Mock load_data to return a DataFrame # Mock load_data to return a DataFrame
mock_df = pd.DataFrame({ mock_df = pd.DataFrame(
'var1': [1, 2, 3, 4, 5], {
'var2': [2, 3, 4, 5, 6], 'var1': [1, 2, 3, 4, 5],
'var3': [3, 4, 5, 6, 7], 'var2': [2, 3, 4, 5, 6],
'target': [10, 15, 20, 25, 30], 'var3': [3, 4, 5, 6, 7],
}) 'target': [10, 15, 20, 25, 30],
}
)
mock_load_data.return_value = mock_df mock_load_data.return_value = mock_df
# Mock split_train_test to return train/test splits # Mock split_train_test to return train/test splits
@@ -540,12 +542,14 @@ class TestTrain:
"""Test training with scaler enabled.""" """Test training with scaler enabled."""
sample_params.use_scaler = True sample_params.use_scaler = True
mock_df = pd.DataFrame({ mock_df = pd.DataFrame(
'var1': [1, 2, 3, 4, 5], {
'var2': [2, 3, 4, 5, 6], 'var1': [1, 2, 3, 4, 5],
'var3': [3, 4, 5, 6, 7], 'var2': [2, 3, 4, 5, 6],
'target': [10, 15, 20, 25, 30], 'var3': [3, 4, 5, 6, 7],
}) 'target': [10, 15, 20, 25, 30],
}
)
mock_load_data.return_value = mock_df mock_load_data.return_value = mock_df
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
@@ -567,12 +571,14 @@ class TestTrain:
"""Test training with shuffle enabled.""" """Test training with shuffle enabled."""
sample_params.shuffle = True sample_params.shuffle = True
mock_df = pd.DataFrame({ mock_df = pd.DataFrame(
'var1': [1, 2, 3, 4, 5], {
'var2': [2, 3, 4, 5, 6], 'var1': [1, 2, 3, 4, 5],
'var3': [3, 4, 5, 6, 7], 'var2': [2, 3, 4, 5, 6],
'target': [10, 15, 20, 25, 30], 'var3': [3, 4, 5, 6, 7],
}) 'target': [10, 15, 20, 25, 30],
}
)
mock_load_data.return_value = mock_df mock_load_data.return_value = mock_df
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
@@ -581,7 +587,7 @@ class TestTrain:
y_test = pd.Series([25, 30], name='target') y_test = pd.Series([25, 30], name='target')
mock_split_train_test.return_value = (x_train, x_test, y_train, y_test) mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)
result = training_repo.train(sample_csv_data, sample_params) training_repo.train(sample_csv_data, sample_params)
# Verify split was called with shuffle=True # Verify split was called with shuffle=True
call_kwargs = mock_split_train_test.call_args[1] call_kwargs = mock_split_train_test.call_args[1]
@@ -595,12 +601,14 @@ class TestTrain:
"""Test training with different train size.""" """Test training with different train size."""
sample_params.train_size = 70 sample_params.train_size = 70
mock_df = pd.DataFrame({ mock_df = pd.DataFrame(
'var1': [1, 2, 3, 4, 5], {
'var2': [2, 3, 4, 5, 6], 'var1': [1, 2, 3, 4, 5],
'var3': [3, 4, 5, 6, 7], 'var2': [2, 3, 4, 5, 6],
'target': [10, 15, 20, 25, 30], 'var3': [3, 4, 5, 6, 7],
}) 'target': [10, 15, 20, 25, 30],
}
)
mock_load_data.return_value = mock_df mock_load_data.return_value = mock_df
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
@@ -609,7 +617,7 @@ class TestTrain:
y_test = pd.Series([25, 30], name='target') y_test = pd.Series([25, 30], name='target')
mock_split_train_test.return_value = (x_train, x_test, y_train, y_test) mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)
result = training_repo.train(sample_csv_data, sample_params) training_repo.train(sample_csv_data, sample_params)
# Verify split was called with train_size=0.7 # Verify split was called with train_size=0.7
call_kwargs = mock_split_train_test.call_args[1] call_kwargs = mock_split_train_test.call_args[1]
@@ -622,12 +630,14 @@ class TestTrain:
): ):
"""Test that ValueError is raised when transformed data is empty.""" """Test that ValueError is raised when transformed data is empty."""
# Mock load_data to return empty DataFrame # Mock load_data to return empty DataFrame
mock_df = pd.DataFrame({ mock_df = pd.DataFrame(
'var1': [], {
'var2': [], 'var1': [],
'var3': [], 'var2': [],
'target': [], 'var3': [],
}) 'target': [],
}
)
mock_load_data.return_value = mock_df mock_load_data.return_value = mock_df
with pytest.raises(ValueError, match='Data view is empty after transformation'): with pytest.raises(ValueError, match='Data view is empty after transformation'):
@@ -645,12 +655,14 @@ class TestTrain:
sample_csv_data, sample_csv_data,
): ):
"""Test that training success is logged.""" """Test that training success is logged."""
mock_df = pd.DataFrame({ mock_df = pd.DataFrame(
'var1': [1, 2, 3, 4, 5], {
'var2': [2, 3, 4, 5, 6], 'var1': [1, 2, 3, 4, 5],
'var3': [3, 4, 5, 6, 7], 'var2': [2, 3, 4, 5, 6],
'target': [10, 15, 20, 25, 30], 'var3': [3, 4, 5, 6, 7],
}) 'target': [10, 15, 20, 25, 30],
}
)
mock_load_data.return_value = mock_df mock_load_data.return_value = mock_df
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
@@ -676,12 +688,14 @@ class TestTrain:
sample_params.line_separator = ';' sample_params.line_separator = ';'
sample_params.decimal_separator = ',' sample_params.decimal_separator = ','
mock_df = pd.DataFrame({ mock_df = pd.DataFrame(
'var1': [1, 2, 3, 4, 5], {
'var2': [2, 3, 4, 5, 6], 'var1': [1, 2, 3, 4, 5],
'var3': [3, 4, 5, 6, 7], 'var2': [2, 3, 4, 5, 6],
'target': [10, 15, 20, 25, 30], 'var3': [3, 4, 5, 6, 7],
}) 'target': [10, 15, 20, 25, 30],
}
)
mock_load_data.return_value = mock_df mock_load_data.return_value = mock_df
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
@@ -701,12 +715,14 @@ class TestTrain:
self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
): ):
"""Test that TrainModelResult contains all expected fields.""" """Test that TrainModelResult contains all expected fields."""
mock_df = pd.DataFrame({ mock_df = pd.DataFrame(
'var1': [1, 2, 3, 4, 5], {
'var2': [2, 3, 4, 5, 6], 'var1': [1, 2, 3, 4, 5],
'var3': [3, 4, 5, 6, 7], 'var2': [2, 3, 4, 5, 6],
'target': [10, 15, 20, 25, 30], 'var3': [3, 4, 5, 6, 7],
}) 'target': [10, 15, 20, 25, 30],
}
)
mock_load_data.return_value = mock_df mock_load_data.return_value = mock_df
x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})

47
tests/worker/test_worker.py
View File

@@ -3,7 +3,7 @@
import asyncio import asyncio
import os import os
import sys import sys
from unittest.mock import AsyncMock, MagicMock, Mock, patch from unittest.mock import AsyncMock, Mock, patch
import pytest import pytest
@@ -144,7 +144,9 @@ def test_start_prometheus_server_custom_port(mock_metrics, mock_start_http_serve
@patch('model_manager.worker.worker.start_http_server') @patch('model_manager.worker.worker.start_http_server')
@patch('model_manager.worker.worker.metrics') @patch('model_manager.worker.worker.metrics')
@patch('model_manager.worker.worker.os._exit') @patch('model_manager.worker.worker.os._exit')
def test_start_prometheus_server_failure(mock_exit, mock_metrics, mock_start_http_server, mock_env_vars): def test_start_prometheus_server_failure(
mock_exit, mock_metrics, mock_start_http_server, mock_env_vars
):
"""Test Prometheus server startup failure.""" """Test Prometheus server startup failure."""
from model_manager.worker.worker import start_prometheus_server from model_manager.worker.worker import start_prometheus_server
@@ -194,7 +196,10 @@ async def test_main_successful_startup(
# Setup mocks # Setup mocks
mock_get_logger.return_value = mock_logger mock_get_logger.return_value = mock_logger
mock_build_mongodb.return_value = {'connection_string': 'mongodb://test', 'database_name': 'test_db'} mock_build_mongodb.return_value = {
'connection_string': 'mongodb://test',
'database_name': 'test_db',
}
mock_build_postgres.return_value = {} mock_build_postgres.return_value = {}
mock_build_mlflow.return_value = {} mock_build_mlflow.return_value = {}
mock_build_minio.return_value = {} mock_build_minio.return_value = {}
@@ -209,7 +214,9 @@ async def test_main_successful_startup(
mock_client_class.connect = AsyncMock(return_value=mock_client_instance) mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_worker_instance = Mock() mock_worker_instance = Mock()
mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError()) # Simulate interruption mock_worker_instance.run = AsyncMock(
side_effect=asyncio.CancelledError()
) # Simulate interruption
mock_worker_class.return_value = mock_worker_instance mock_worker_class.return_value = mock_worker_instance
mock_app_up = Mock() mock_app_up = Mock()
@@ -269,7 +276,10 @@ async def test_main_handles_exception(
# Setup mocks # Setup mocks
mock_get_logger.return_value = mock_logger mock_get_logger.return_value = mock_logger
mock_build_mongodb.return_value = {'connection_string': 'mongodb://test', 'database_name': 'test_db'} mock_build_mongodb.return_value = {
'connection_string': 'mongodb://test',
'database_name': 'test_db',
}
mock_build_postgres.return_value = {} mock_build_postgres.return_value = {}
mock_build_mlflow.return_value = {} mock_build_mlflow.return_value = {}
mock_build_minio.return_value = {} mock_build_minio.return_value = {}
@@ -342,10 +352,16 @@ async def test_main_temporal_client_configuration(
"""Test that Temporal client is configured correctly.""" """Test that Temporal client is configured correctly."""
from model_manager.worker.worker import main from model_manager.worker.worker import main
with patch.dict(os.environ, {'TEMPORAL_HOST': 'temporal.example.com:7233', 'TEMPORAL_NAMESPACE': 'production'}): with patch.dict(
os.environ,
{'TEMPORAL_HOST': 'temporal.example.com:7233', 'TEMPORAL_NAMESPACE': 'production'},
):
# Setup mocks # Setup mocks
mock_get_logger.return_value = mock_logger mock_get_logger.return_value = mock_logger
mock_build_mongodb.return_value = {'connection_string': 'mongodb://test', 'database_name': 'test_db'} mock_build_mongodb.return_value = {
'connection_string': 'mongodb://test',
'database_name': 'test_db',
}
mock_build_postgres.return_value = {} mock_build_postgres.return_value = {}
mock_build_mlflow.return_value = {} mock_build_mlflow.return_value = {}
mock_build_minio.return_value = {} mock_build_minio.return_value = {}
@@ -415,7 +431,10 @@ async def test_main_worker_configuration(
# Setup mocks # Setup mocks
mock_get_logger.return_value = mock_logger mock_get_logger.return_value = mock_logger
mock_build_mongodb.return_value = {'connection_string': 'mongodb://test', 'database_name': 'test_db'} mock_build_mongodb.return_value = {
'connection_string': 'mongodb://test',
'database_name': 'test_db',
}
mock_build_postgres.return_value = {} mock_build_postgres.return_value = {}
mock_build_mlflow.return_value = {} mock_build_mlflow.return_value = {}
mock_build_minio.return_value = {} mock_build_minio.return_value = {}
@@ -472,8 +491,6 @@ def test_main_entrypoint(mock_asyncio_run):
"""Test the __main__ entrypoint.""" """Test the __main__ entrypoint."""
# Import and execute the main block # Import and execute the main block
with patch.object(sys, 'argv', ['worker.py']): with patch.object(sys, 'argv', ['worker.py']):
import importlib
import model_manager.worker.worker as worker_module import model_manager.worker.worker as worker_module
# Simulate running the module # Simulate running the module
@@ -495,7 +512,9 @@ def test_worker_module_docstring():
@patch('model_manager.worker.worker.start_http_server') @patch('model_manager.worker.worker.start_http_server')
@patch('model_manager.worker.worker.metrics') @patch('model_manager.worker.worker.metrics')
def test_start_prometheus_server_prints_success(mock_metrics, mock_start_http_server, capsys, mock_env_vars): def test_start_prometheus_server_prints_success(
mock_metrics, mock_start_http_server, capsys, mock_env_vars
):
"""Test that start_prometheus_server prints success message.""" """Test that start_prometheus_server prints success message."""
from model_manager.worker.worker import start_prometheus_server from model_manager.worker.worker import start_prometheus_server
@@ -511,7 +530,9 @@ def test_start_prometheus_server_prints_success(mock_metrics, mock_start_http_se
@patch('model_manager.worker.worker.start_http_server') @patch('model_manager.worker.worker.start_http_server')
@patch('model_manager.worker.worker.metrics') @patch('model_manager.worker.worker.metrics')
@patch('model_manager.worker.worker.os._exit') @patch('model_manager.worker.worker.os._exit')
def test_start_prometheus_server_prints_failure(mock_exit, mock_metrics, mock_start_http_server, capsys, mock_env_vars): def test_start_prometheus_server_prints_failure(
mock_exit, mock_metrics, mock_start_http_server, capsys, mock_env_vars
):
"""Test that start_prometheus_server prints failure message.""" """Test that start_prometheus_server prints failure message."""
from model_manager.worker.worker import start_prometheus_server from model_manager.worker.worker import start_prometheus_server
@@ -522,5 +543,3 @@ def test_start_prometheus_server_prints_failure(mock_exit, mock_metrics, mock_st
captured = capsys.readouterr() captured = capsys.readouterr()
assert 'Failed to start Prometheus server' in captured.out assert 'Failed to start Prometheus server' in captured.out
assert 'Test error' in captured.out assert 'Test error' in captured.out

104
tests/workflows/test_train_model.py
View File

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