SIENTIAPDE-1241: refactor train_model workflow due to I/O errors.

This commit is contained in:
Bruno Domingues
2025-10-22 15:37:56 -03:00
parent f2a1c88ff3
commit 5789a13023
31 changed files with 37878 additions and 6097 deletions

View File

@@ -1,708 +0,0 @@
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
from pandas import DataFrame
from model_manager.utils.repository.model_repository import MLFlowRepository
@pytest.fixture
def mlflow_repository():
with patch(
'model_manager.utils.repository.model_repository.ModelServing', autospec=True
) as mock_model_serving:
mock_instance = mock_model_serving.return_value
mock_instance.get_transformed_data = MagicMock()
repo = MLFlowRepository(
host='http://localhost:5000', username='admin', password='admin', logger=MagicMock()
)
return repo
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
# ========== Tests for Model Artifact Generation Methods ==========
def test_get_next_run_name(mlflow_repository):
"""Test get_next_run_name generates correct run name based on existing runs."""
mlflow_repository.model_serving.search_runs_by_name.return_value = [
MagicMock(),
MagicMock(),
MagicMock(),
]
result = mlflow_repository.get_next_run_name('test_experiment')
mlflow_repository.model_serving.search_runs_by_name.assert_called_once_with(
experiment_names=['test_experiment'], order_by=['start_time desc']
)
assert result == 'test_experiment-4'
def test_get_next_run_name_first_run(mlflow_repository):
"""Test get_next_run_name for first run (no existing runs)."""
mlflow_repository.model_serving.search_runs_by_name.return_value = []
result = mlflow_repository.get_next_run_name('test_experiment')
assert result == 'test_experiment-1'
@patch('model_manager.utils.repository.model_repository.path')
def test_generate_artifacts_success(mock_path, mlflow_repository):
"""Test generate_artifacts successfully creates all artifacts."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
# Mock data
params = MagicMock(spec=TrainModelParams)
params.target_variable = 'target'
params.variable_columns = ['feat1', 'feat2']
params.experiment_name = 'test_exp'
data = MagicMock(spec=TrainModelResult)
data.run_name = 'test_run-1'
data.params = params
data.x_train = DataFrame({'feat1': [1, 2], 'feat2': [3, 4]})
data.y_train = DataFrame({'target': [5, 6]})
data.x_test = DataFrame({'feat1': [7, 8], 'feat2': [9, 10]})
data.y_test = DataFrame({'target': [11, 12]})
data.regr = MagicMock()
data.regr.predict = MagicMock(return_value=np.array([5.1, 6.1]))
data.y_pred = np.array([11.1, 12.1])
# Mock path operations
mock_path.exists.return_value = True
mock_path.join.side_effect = lambda *args: '/'.join(args)
# Mock private methods
mlflow_repository._get_reports_directory = MagicMock(return_value='/reports')
mlflow_repository._create_run_directory = MagicMock(return_value='/reports/test_run-1_20231010')
mlflow_repository._setup_run_directory = MagicMock()
mlflow_repository._generate_report = MagicMock(return_value=data)
result = mlflow_repository.generate_artifacts(data)
# Assertions
mlflow_repository._get_reports_directory.assert_called_once()
mlflow_repository._create_run_directory.assert_called_once_with('/reports', 'test_run-1')
mlflow_repository._setup_run_directory.assert_called_once()
mlflow_repository._generate_report.assert_called_once()
assert result == data
def test_generate_artifacts_missing_run_name(mlflow_repository):
"""Test generate_artifacts raises ValueError when run_name is not set."""
from model_manager.utils.models.train_model_result import TrainModelResult
data = MagicMock(spec=TrainModelResult)
data.run_name = None
with pytest.raises(ValueError) as exc_info:
mlflow_repository.generate_artifacts(data)
assert 'run_name must be set before generating artifacts' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.path')
def test_generate_artifacts_reports_directory_not_exists(mock_path, mlflow_repository):
"""Test generate_artifacts raises FileNotFoundError when reports directory doesn't exist."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
params.target_variable = 'target'
data = MagicMock(spec=TrainModelResult)
data.run_name = 'test_run-1'
data.params = params
data.x_train = DataFrame({'feat1': [1]})
data.y_train = DataFrame({'target': [2]})
data.x_test = DataFrame({'feat1': [3]})
data.y_test = DataFrame({'target': [4]})
data.regr = MagicMock()
data.y_pred = np.array([4.1])
mlflow_repository._get_reports_directory = MagicMock(return_value='/reports')
mock_path.exists.return_value = False
with pytest.raises(FileNotFoundError) as exc_info:
mlflow_repository.generate_artifacts(data)
assert 'Reports directory does not exist' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.path')
def test_generate_artifacts_header_file_not_exists(mock_path, mlflow_repository):
"""Test generate_artifacts raises FileNotFoundError when header.html doesn't exist."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
params.target_variable = 'target'
data = MagicMock(spec=TrainModelResult)
data.run_name = 'test_run-1'
data.params = params
data.x_train = DataFrame({'feat1': [1]})
data.y_train = DataFrame({'target': [2]})
data.x_test = DataFrame({'feat1': [3]})
data.y_test = DataFrame({'target': [4]})
data.regr = MagicMock()
data.regr.predict = MagicMock(return_value=np.array([2.1]))
data.y_pred = np.array([4.1])
mlflow_repository._get_reports_directory = MagicMock(return_value='/reports')
mlflow_repository._create_run_directory = MagicMock(return_value='/reports/test_run-1_20231010')
# First call returns True (reports dir exists), second returns False (header.html doesn't exist)
mock_path.exists.side_effect = [True, False]
mock_path.join.side_effect = lambda *args: '/'.join(args)
with pytest.raises(FileNotFoundError) as exc_info:
mlflow_repository.generate_artifacts(data)
assert 'Header file does not exist' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.path')
def test_save_run_success(mock_path, mlflow_repository):
"""Test save_run successfully logs all parameters, metrics, models, and artifacts."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
params.train_size = 80
params.removed_intervals = [(1, 10), (20, 30)]
params.experiment_name = 'test_exp'
params.target_variable = 'target'
params.variable_columns = ['feat1', 'feat2']
params.lag_train = 5
params.lag_val = 3
params.window = 10
params.low_lim = 0.0
params.upp_lim = 1.0
params.include_ar = True
data = MagicMock(spec=TrainModelResult)
data.run_name = 'test_run-1'
data.params = params
data.report_path = '/reports/report.html'
data.train_data_path = '/reports/train.csv'
data.test_data_path = '/reports/test.csv'
data.mse_val = 0.123
data.r2_val = 0.987
data.mae_val = 0.456
data.scaler_dict = {'scaler': 'minmax'}
data.process_data = MagicMock()
data.regr = MagicMock()
mock_path.exists.return_value = True
mlflow_repository.save_run(data)
# Verify experiment was set
mlflow_repository.model_serving.set_experiment.assert_called_once_with('test_exp')
# Verify parameters were logged
assert mlflow_repository.model_serving.log_param.call_count == 13
# Verify metrics were logged
mlflow_repository.model_serving.log_metric.assert_any_call('MSE', 0.123)
mlflow_repository.model_serving.log_metric.assert_any_call('R2', 0.987)
mlflow_repository.model_serving.log_metric.assert_any_call('MAE', 0.456)
# Verify models were logged
mlflow_repository.model_serving.log_model.assert_any_call(data.process_data, 'data_model')
mlflow_repository.model_serving.log_model.assert_any_call(data.regr, 'prediction_model')
# Verify artifacts were logged
mlflow_repository.model_serving.log_artifact.assert_any_call('/reports/report.html')
mlflow_repository.model_serving.log_artifact.assert_any_call('/reports/train.csv')
mlflow_repository.model_serving.log_artifact.assert_any_call('/reports/test.csv')
@patch('model_manager.utils.repository.model_repository.path')
def test_save_run_missing_report_path(mock_path, mlflow_repository):
"""Test save_run raises ValueError when report_path is missing."""
from model_manager.utils.models.train_model_result import TrainModelResult
data = MagicMock(spec=TrainModelResult)
data.report_path = None
with pytest.raises(ValueError) as exc_info:
mlflow_repository.save_run(data)
assert 'Report file does not exist' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.path')
def test_save_run_missing_metrics(mock_path, mlflow_repository):
"""Test save_run raises ValueError when metrics are None."""
from model_manager.utils.models.train_model_result import TrainModelResult
data = MagicMock(spec=TrainModelResult)
data.report_path = '/reports/report.html'
data.train_data_path = '/reports/train.csv'
data.test_data_path = '/reports/test.csv'
data.mse_val = None
data.r2_val = 0.987
data.mae_val = 0.456
mock_path.exists.return_value = True
with pytest.raises(ValueError) as exc_info:
mlflow_repository.save_run(data)
assert 'One or more metrics (MSE, R2, MAE) are None' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.path')
def test_save_run_mlflow_error(mock_path, mlflow_repository):
"""Test save_run handles MLflow errors gracefully."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
params.train_size = 80
params.removed_intervals = []
params.experiment_name = 'test_exp'
data = MagicMock(spec=TrainModelResult)
data.run_name = 'test_run-1'
data.params = params
data.report_path = '/reports/report.html'
data.train_data_path = '/reports/train.csv'
data.test_data_path = '/reports/test.csv'
data.mse_val = 0.123
data.r2_val = 0.987
data.mae_val = 0.456
mock_path.exists.return_value = True
mlflow_repository.model_serving.set_experiment.side_effect = Exception(
'MLflow connection error'
)
with pytest.raises(RuntimeError) as exc_info:
mlflow_repository.save_run(data)
assert 'Failed to save run' in str(exc_info.value)
assert 'MLflow connection error' in str(exc_info.value)
# ========== Additional Tests for 100% Coverage ==========
def test_init_artifacts_data_empty_x_train(mlflow_repository):
"""Test _init_artifacts_data raises ValueError when x_train is empty."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
data = MagicMock(spec=TrainModelResult)
data.params = params
data.x_train = DataFrame() # Empty DataFrame
data.y_train = DataFrame({'target': [1]})
data.x_test = DataFrame({'feat1': [1]})
data.y_test = DataFrame({'target': [1]})
with pytest.raises(ValueError) as exc_info:
mlflow_repository._init_artifacts_data(data)
assert 'Training features (x_train) are empty' in str(exc_info.value)
def test_init_artifacts_data_empty_y_train(mlflow_repository):
"""Test _init_artifacts_data raises ValueError when y_train is empty."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
data = MagicMock(spec=TrainModelResult)
data.params = params
data.x_train = DataFrame({'feat1': [1]})
data.y_train = DataFrame() # Empty DataFrame
data.x_test = DataFrame({'feat1': [1]})
data.y_test = DataFrame({'target': [1]})
with pytest.raises(ValueError) as exc_info:
mlflow_repository._init_artifacts_data(data)
assert 'Training target (y_train) is empty' in str(exc_info.value)
def test_init_artifacts_data_empty_x_test(mlflow_repository):
"""Test _init_artifacts_data raises ValueError when x_test is empty."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
data = MagicMock(spec=TrainModelResult)
data.params = params
data.x_train = DataFrame({'feat1': [1]})
data.y_train = DataFrame({'target': [1]})
data.x_test = DataFrame() # Empty DataFrame
data.y_test = DataFrame({'target': [1]})
with pytest.raises(ValueError) as exc_info:
mlflow_repository._init_artifacts_data(data)
assert 'Test features (x_test) are empty' in str(exc_info.value)
def test_init_artifacts_data_empty_y_test(mlflow_repository):
"""Test _init_artifacts_data raises ValueError when y_test is empty."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
data = MagicMock(spec=TrainModelResult)
data.params = params
data.x_train = DataFrame({'feat1': [1]})
data.y_train = DataFrame({'target': [1]})
data.x_test = DataFrame({'feat1': [1]})
data.y_test = DataFrame() # Empty DataFrame
with pytest.raises(ValueError) as exc_info:
mlflow_repository._init_artifacts_data(data)
assert 'Test target (y_test) is empty' in str(exc_info.value)
def test_init_artifacts_data_none_y_pred(mlflow_repository):
"""Test _init_artifacts_data raises ValueError when y_pred is None."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
data = MagicMock(spec=TrainModelResult)
data.params = params
data.x_train = DataFrame({'feat1': [1]})
data.y_train = DataFrame({'target': [1]})
data.x_test = DataFrame({'feat1': [1]})
data.y_test = DataFrame({'target': [1]})
data.y_pred = None
with pytest.raises(ValueError) as exc_info:
mlflow_repository._init_artifacts_data(data)
assert 'Test predictions (y_pred) are None' in str(exc_info.value)
def test_init_artifacts_data_success(mlflow_repository):
"""Test _init_artifacts_data successfully prepares data."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
params.target_variable = 'target'
data = MagicMock(spec=TrainModelResult)
data.params = params
data.x_train = DataFrame({'feat1': [1, 2]})
data.y_train = DataFrame({'target': [3, 4]})
data.x_test = DataFrame({'feat1': [5, 6]})
data.y_test = DataFrame({'target': [7, 8]})
data.regr = MagicMock()
data.regr.predict = MagicMock(return_value=np.array([3.1, 4.1]))
data.y_pred = np.array([7.1, 8.1])
reference_data, current_data = mlflow_repository._init_artifacts_data(data)
assert 'target' in reference_data.columns
assert 'prediction' in reference_data.columns
assert 'target' in current_data.columns
assert 'prediction' in current_data.columns
assert len(reference_data) == 2
assert len(current_data) == 2
@patch('model_manager.utils.repository.model_repository.makedirs')
@patch('model_manager.utils.repository.model_repository.path')
def test_create_run_directory_success(mock_path, mock_makedirs, mlflow_repository):
"""Test _create_run_directory successfully creates directory."""
mock_path.join.return_value = '/reports/test_run_20231010_123456_123456'
result = mlflow_repository._create_run_directory('/reports', 'test_run')
mock_makedirs.assert_called_once_with('/reports/test_run_20231010_123456_123456', exist_ok=True)
assert result == '/reports/test_run_20231010_123456_123456'
@patch('model_manager.utils.repository.model_repository.makedirs')
@patch('model_manager.utils.repository.model_repository.path')
def test_create_run_directory_permission_error(mock_path, mock_makedirs, mlflow_repository):
"""Test _create_run_directory handles PermissionError."""
mock_path.join.return_value = '/reports/test_run_20231010'
mock_makedirs.side_effect = PermissionError('Permission denied')
with pytest.raises(PermissionError) as exc_info:
mlflow_repository._create_run_directory('/reports', 'test_run')
assert 'Permission denied when creating directory' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.makedirs')
@patch('model_manager.utils.repository.model_repository.path')
def test_create_run_directory_os_error(mock_path, mock_makedirs, mlflow_repository):
"""Test _create_run_directory handles OSError."""
mock_path.join.return_value = '/reports/test_run_20231010'
mock_makedirs.side_effect = OSError('Disk full')
with pytest.raises(OSError) as exc_info:
mlflow_repository._create_run_directory('/reports', 'test_run')
assert 'Failed to create directory' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.shutil')
@patch('model_manager.utils.repository.model_repository.path')
def test_setup_run_directory_success(mock_path, mock_shutil, mlflow_repository):
"""Test _setup_run_directory successfully sets up directory."""
mock_path.join.side_effect = lambda *args: '/'.join(args)
mock_open = MagicMock()
with patch('builtins.open', mock_open):
mlflow_repository._setup_run_directory('/run_dir', '/reports/header.html')
assert mock_open.call_count == 3 # 3 empty files
mock_shutil.copy.assert_called_once_with('/reports/header.html', '/run_dir/header.html')
@patch('model_manager.utils.repository.model_repository.shutil')
@patch('model_manager.utils.repository.model_repository.path')
def test_setup_run_directory_file_not_found(mock_path, mock_shutil, mlflow_repository):
"""Test _setup_run_directory handles FileNotFoundError."""
mock_path.join.side_effect = lambda *args: '/'.join(args)
mock_shutil.copy.side_effect = FileNotFoundError('Header not found')
mock_open = MagicMock()
with patch('builtins.open', mock_open):
with pytest.raises(FileNotFoundError) as exc_info:
mlflow_repository._setup_run_directory('/run_dir', '/reports/header.html')
assert 'Header file not found' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.shutil')
@patch('model_manager.utils.repository.model_repository.path')
def test_setup_run_directory_permission_error(mock_path, mock_shutil, mlflow_repository):
"""Test _setup_run_directory handles PermissionError."""
mock_path.join.side_effect = lambda *args: '/'.join(args)
mock_open = MagicMock()
mock_open.side_effect = PermissionError('Permission denied')
with patch('builtins.open', mock_open):
with pytest.raises(PermissionError) as exc_info:
mlflow_repository._setup_run_directory('/run_dir', '/reports/header.html')
assert 'Permission denied when setting up directory' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.shutil')
@patch('model_manager.utils.repository.model_repository.path')
def test_setup_run_directory_os_error(mock_path, mock_shutil, mlflow_repository):
"""Test _setup_run_directory handles OSError."""
mock_path.join.side_effect = lambda *args: '/'.join(args)
mock_open = MagicMock()
mock_open.side_effect = OSError('Disk error')
with patch('builtins.open', mock_open):
with pytest.raises(OSError) as exc_info:
mlflow_repository._setup_run_directory('/run_dir', '/reports/header.html')
assert 'Failed to setup run directory' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.Reports')
@patch('model_manager.utils.repository.model_repository.path')
def test_generate_report_success(mock_path, mock_reports, mlflow_repository):
"""Test _generate_report successfully generates all reports."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
params.variable_columns = ['feat1', 'feat2']
data = MagicMock(spec=TrainModelResult)
data.params = params
data.run_dir = '/run_dir'
reference_data = DataFrame(
{'feat1': [1.0], 'feat2': [2.0], 'target': [3.0], 'prediction': [3.1]}
)
current_data = DataFrame({'feat1': [4.0], 'feat2': [5.0], 'target': [6.0], 'prediction': [6.1]})
mock_path.join.side_effect = lambda *args: '/'.join(args)
mock_report_instance = MagicMock()
mock_reports.return_value = mock_report_instance
# Mock DataFrame.to_csv to avoid actual file writing
with patch.object(DataFrame, 'to_csv'):
result = mlflow_repository._generate_report(reference_data, current_data, data)
mock_reports.assert_called_once()
mock_report_instance.add_data_quality_section.assert_called_once()
mock_report_instance.add_data_drift_section.assert_called_once()
mock_report_instance.add_regression_section.assert_called_once()
mock_report_instance.save_all_sections_html.assert_called_once()
assert result == data
@patch('model_manager.utils.repository.model_repository.path')
def test_generate_report_value_error(mock_path, mlflow_repository):
"""Test _generate_report handles ValueError from data conversion."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
data = MagicMock(spec=TrainModelResult)
data.params = params
data.run_dir = '/run_dir'
# DataFrame with non-numeric data
reference_data = DataFrame({'feat1': ['a', 'b']})
current_data = DataFrame({'feat1': ['c', 'd']})
with pytest.raises(ValueError) as exc_info:
mlflow_repository._generate_report(reference_data, current_data, data)
assert 'Failed to convert data to float64' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.Reports')
@patch('model_manager.utils.repository.model_repository.path')
def test_generate_report_permission_error(mock_path, mock_reports, mlflow_repository):
"""Test _generate_report handles PermissionError."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
params.variable_columns = ['feat1']
data = MagicMock(spec=TrainModelResult)
data.params = params
data.run_dir = '/run_dir'
reference_data = DataFrame({'feat1': [1.0]})
current_data = DataFrame({'feat1': [2.0]})
mock_path.join.side_effect = lambda *args: '/'.join(args)
mock_report_instance = MagicMock()
mock_reports.return_value = mock_report_instance
mock_report_instance.save_all_sections_html.side_effect = PermissionError('Permission denied')
with pytest.raises(PermissionError) as exc_info:
mlflow_repository._generate_report(reference_data, current_data, data)
assert 'Permission denied when writing report files' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.Reports')
@patch('model_manager.utils.repository.model_repository.path')
def test_generate_report_os_error(mock_path, mock_reports, mlflow_repository):
"""Test _generate_report handles OSError."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
params.variable_columns = ['feat1']
data = MagicMock(spec=TrainModelResult)
data.params = params
data.run_dir = '/run_dir'
reference_data = DataFrame({'feat1': [1.0]})
current_data = DataFrame({'feat1': [2.0]})
mock_path.join.side_effect = lambda *args: '/'.join(args)
mock_report_instance = MagicMock()
mock_reports.return_value = mock_report_instance
mock_report_instance.save_all_sections_html.side_effect = OSError('Disk error')
with pytest.raises(OSError) as exc_info:
mlflow_repository._generate_report(reference_data, current_data, data)
assert 'Failed to generate report' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.Reports')
@patch('model_manager.utils.repository.model_repository.path')
def test_generate_report_run_dir_none(mock_path, mock_reports, mlflow_repository):
"""Test _generate_report raises ValueError when run_dir is None."""
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
params = MagicMock(spec=TrainModelParams)
params.variable_columns = ['feat1']
data = MagicMock(spec=TrainModelResult)
data.params = params
data.run_dir = None # Not set
reference_data = DataFrame({'feat1': [1.0]})
current_data = DataFrame({'feat1': [2.0]})
mock_report_instance = MagicMock()
mock_reports.return_value = mock_report_instance
with pytest.raises(ValueError) as exc_info:
mlflow_repository._generate_report(reference_data, current_data, data)
assert 'run_dir is not set after directory creation' in str(exc_info.value)
def test_get_reports_directory(mlflow_repository):
"""Test _get_reports_directory returns correct path."""
result = mlflow_repository._get_reports_directory()
assert result.endswith('model_manager/reports')
assert 'model_manager' in result
@patch('model_manager.utils.repository.model_repository.path')
def test_save_run_missing_train_data_path(mock_path, mlflow_repository):
"""Test save_run raises ValueError when train_data_path is missing."""
from model_manager.utils.models.train_model_result import TrainModelResult
data = MagicMock(spec=TrainModelResult)
data.report_path = '/reports/report.html'
data.train_data_path = None
mock_path.exists.return_value = True
with pytest.raises(ValueError) as exc_info:
mlflow_repository.save_run(data)
assert 'Training data file does not exist' in str(exc_info.value)
@patch('model_manager.utils.repository.model_repository.path')
def test_save_run_missing_test_data_path(mock_path, mlflow_repository):
"""Test save_run raises ValueError when test_data_path is missing."""
from model_manager.utils.models.train_model_result import TrainModelResult
data = MagicMock(spec=TrainModelResult)
data.report_path = '/reports/report.html'
data.train_data_path = '/reports/train.csv'
data.test_data_path = None
mock_path.exists.return_value = True
with pytest.raises(ValueError) as exc_info:
mlflow_repository.save_run(data)
assert 'Test data file does not exist' in str(exc_info.value)

View File

@@ -1,324 +0,0 @@
"""Unit tests for TrainingRepository."""
from io import BytesIO
from unittest.mock import MagicMock, patch
import numpy as np
import pandas as pd
from pytest import fixture, raises
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
from model_manager.utils.repository.training_repository import TrainingRepository
@fixture
def logger():
"""Create a mock logger."""
return MagicMock()
@fixture
def training_repository(logger):
"""Create a TrainingRepository instance."""
return TrainingRepository(logger)
@fixture
def train_params():
"""Create sample training parameters."""
return TrainModelParams(
variable_columns=['feature1', 'feature2'],
lag_train=1,
lag_val=1,
target_variable='target',
rem_static_win=False,
low_lim={'feature1': 0.0, 'feature2': 0.0},
upp_lim={'feature1': 100.0, 'feature2': 100.0},
window=10,
use_scaler=True,
include_ar=False,
bucket_name='test-bucket',
file_name='test.csv',
line_separator='\n',
decimal_separator='.',
train_size=80,
shuffle=True,
experiment_run_id=123,
experiment_name='test_experiment',
removed_intervals=[],
)
@fixture
def sample_csv_data():
"""Create sample CSV data."""
csv_content = """feature1,feature2,target
1.0,2.0,10.0
2.0,3.0,15.0
3.0,4.0,20.0
4.0,5.0,25.0
5.0,6.0,30.0
6.0,7.0,35.0
7.0,8.0,40.0
8.0,9.0,45.0
9.0,10.0,50.0
10.0,11.0,55.0
"""
return BytesIO(csv_content.encode())
@patch('model_manager.utils.repository.training_repository.load_data')
@patch('model_manager.utils.repository.training_repository.DataPreprocessor')
@patch('model_manager.utils.repository.training_repository.split_train_test')
@patch('model_manager.utils.repository.training_repository.LinearRegressionModel')
def test_train_success(
mock_linear_model,
mock_split,
mock_preprocessor_class,
mock_load_data,
training_repository,
train_params,
sample_csv_data,
):
"""Test successful model training."""
# Setup mocks
mock_data = pd.DataFrame(
{'feature1': [1, 2, 3, 4, 5], 'feature2': [2, 3, 4, 5, 6], 'target': [10, 15, 20, 25, 30]}
)
mock_load_data.return_value = mock_data
mock_preprocessor = MagicMock()
mock_preprocessor_class.return_value = mock_preprocessor
mock_preprocessor.transform.return_value = mock_data
x_train = pd.DataFrame({'feature1': [1, 2, 3], 'feature2': [2, 3, 4]})
x_test = pd.DataFrame({'feature1': [4, 5], 'feature2': [5, 6]})
y_train = pd.Series([10, 15, 20], name='target')
y_test = pd.Series([25, 30], name='target')
mock_split.return_value = (x_train, x_test, y_train, y_test)
mock_model = MagicMock()
mock_linear_model.return_value = mock_model
mock_scaler = MagicMock()
mock_preprocessor.get_scaler.return_value = mock_scaler
# Execute
result = training_repository.train(sample_csv_data, train_params)
# Assertions
assert isinstance(result, TrainModelResult)
assert result.params == train_params
assert result.process_data == mock_preprocessor
assert result.regr == mock_model
mock_load_data.assert_called_once()
mock_preprocessor.fit.assert_called_once()
mock_model.fit.assert_called_once()
@patch('model_manager.utils.repository.training_repository.load_data')
def test_train_empty_data_after_transform(
mock_load_data, training_repository, train_params, sample_csv_data
):
"""Test training with empty data after transformation."""
mock_data = pd.DataFrame({'feature1': [], 'feature2': [], 'target': []})
mock_load_data.return_value = mock_data
with patch.object(training_repository, 'init_data_preprocessor') as mock_init:
mock_preprocessor = MagicMock()
mock_init.return_value = mock_preprocessor
mock_preprocessor.transform.return_value = pd.DataFrame()
with raises(ValueError, match='Data view is empty after transformation'):
training_repository.train(sample_csv_data, train_params)
def test_init_scaler_dict_with_minmax_scaler(training_repository, train_params):
"""Test scaler dict initialization with MinMaxScaler."""
mock_preprocessor = MagicMock()
mock_scaler = MagicMock()
mock_scaler.x_min = [0.0, 1.0]
mock_scaler.x_max = [10.0, 11.0]
mock_scaler.y_min = 5.0
mock_scaler.y_max = 50.0
mock_preprocessor.get_scaler.return_value = mock_scaler
# Patch isinstance to return True for MinMaxScaler
with patch(
'model_manager.utils.repository.training_repository.isinstance',
side_effect=lambda obj, cls: cls.__name__ == 'MinMaxScaler',
):
result = training_repository.init_scaler_dict(mock_preprocessor, train_params)
assert result is not None
assert 'feature1' in result
assert 'feature2' in result
assert 'target' in result
assert result['feature1'] == {'min': 0.0, 'max': 10.0}
assert result['feature2'] == {'min': 1.0, 'max': 11.0}
assert result['target'] == {'min': 5.0, 'max': 50.0}
def test_init_scaler_dict_with_z_scaler(training_repository, train_params):
"""Test scaler dict initialization with Z_Scaler."""
mock_preprocessor = MagicMock()
mock_scaler = MagicMock()
mock_scaler.create_dict.return_value = {'mean': 5.0, 'std': 2.0}
mock_preprocessor.get_scaler.return_value = mock_scaler
# Patch isinstance to return True for Z_Scaler
with patch(
'model_manager.utils.repository.training_repository.isinstance',
side_effect=lambda obj, cls: cls.__name__ == 'Z_Scaler',
):
result = training_repository.init_scaler_dict(mock_preprocessor, train_params)
assert result == {'mean': 5.0, 'std': 2.0}
mock_scaler.create_dict.assert_called_once()
def test_init_scaler_dict_without_scaler(training_repository):
"""Test scaler dict initialization when use_scaler is False."""
train_params_no_scaler = TrainModelParams(
variable_columns=['feature1'],
lag_train=1,
lag_val=1,
target_variable='target',
rem_static_win=False,
low_lim={'feature1': 0.0},
upp_lim={'feature1': 100.0},
window=10,
use_scaler=False,
include_ar=False,
bucket_name='test',
file_name='test.csv',
line_separator='\n',
decimal_separator='.',
train_size=80,
shuffle=True,
experiment_run_id=123,
experiment_name='test',
removed_intervals=[],
)
mock_preprocessor = MagicMock()
result = training_repository.init_scaler_dict(mock_preprocessor, train_params_no_scaler)
assert result == {}
@patch('model_manager.utils.repository.training_repository.mse')
@patch('model_manager.utils.repository.training_repository.mae')
@patch('model_manager.utils.repository.training_repository.r2')
def test_after_train_calculation_with_scaler(
mock_r2, mock_mae, mock_mse, training_repository, train_params
):
"""Test post-training calculations with scaler."""
# Setup mock train result
mock_train_result = MagicMock(spec=TrainModelResult)
mock_train_result.params = train_params
mock_train_result.x_train = pd.DataFrame({'feature1': [1, 2, 3], 'feature2': [2, 3, 4]})
mock_train_result.x_test = pd.DataFrame({'feature1': [4, 5], 'feature2': [5, 6]})
mock_train_result.y_train = pd.Series([10, 15, 20], name='target')
mock_train_result.y_test = pd.Series([25, 30], name='target')
mock_regr = MagicMock()
mock_regr.predict.return_value = np.array([24.5, 29.5])
mock_train_result.regr = mock_regr
mock_scaler = MagicMock()
mock_scaler.denormalize_single_input.side_effect = lambda x, col: x
mock_scaler.denormalize_predictions.side_effect = lambda x, col: x
mock_process_data = MagicMock()
mock_process_data.get_scaler.return_value = mock_scaler
mock_train_result.process_data = mock_process_data
# Setup metric mocks
mock_mse.return_value = 0.5
mock_mae.return_value = 0.3
mock_r2.return_value = 0.95
# Execute
result = training_repository.after_train_calculation(train_params, mock_train_result)
# Assertions
assert result == mock_train_result
assert result.mse_val == 0.5
assert result.mae_val == 0.3
assert result.r2_val == 0.95
assert result.y_pred is not None
mock_regr.predict.assert_called_once()
mock_mse.assert_called_once()
mock_mae.assert_called_once()
mock_r2.assert_called_once()
@patch('model_manager.utils.repository.training_repository.mse')
@patch('model_manager.utils.repository.training_repository.mae')
@patch('model_manager.utils.repository.training_repository.r2')
def test_after_train_calculation_without_scaler(mock_r2, mock_mae, mock_mse, training_repository):
"""Test post-training calculations without scaler."""
train_params_no_scaler = TrainModelParams(
variable_columns=['feature1'],
lag_train=1,
lag_val=1,
target_variable='target',
rem_static_win=False,
low_lim={'feature1': 0.0},
upp_lim={'feature1': 100.0},
window=10,
use_scaler=False,
include_ar=False,
bucket_name='test',
file_name='test.csv',
line_separator='\n',
decimal_separator='.',
train_size=80,
shuffle=True,
experiment_run_id=123,
experiment_name='test',
removed_intervals=[],
)
mock_train_result = MagicMock(spec=TrainModelResult)
mock_train_result.params = train_params_no_scaler
mock_train_result.x_train = pd.DataFrame({'feature1': [1, 2, 3]})
mock_train_result.x_test = pd.DataFrame({'feature1': [4, 5]})
mock_train_result.y_train = pd.Series([10, 15, 20], name='target')
mock_train_result.y_test = pd.Series([25, 30], name='target')
mock_regr = MagicMock()
mock_regr.predict.return_value = np.array([24.5, 29.5])
mock_train_result.regr = mock_regr
# Setup metric mocks
mock_mse.return_value = 0.5
mock_mae.return_value = 0.3
mock_r2.return_value = 0.95
# Execute
result = training_repository.after_train_calculation(train_params_no_scaler, mock_train_result)
# Assertions
assert result.mse_val == 0.5
assert result.mae_val == 0.3
assert result.r2_val == 0.95
@patch('model_manager.utils.repository.training_repository.DataPreprocessor')
def test_init_data_preprocessor(mock_preprocessor_class, training_repository, train_params):
"""Test DataPreprocessor initialization."""
mock_preprocessor = MagicMock()
mock_preprocessor_class.return_value = mock_preprocessor
result = training_repository.init_data_preprocessor(train_params)
assert result == mock_preprocessor
mock_preprocessor_class.assert_called_once()
call_kwargs = mock_preprocessor_class.call_args[1]
assert call_kwargs['target_variable'] == 'target'
assert call_kwargs['input_columns'] == ['feature1', 'feature2']
assert call_kwargs['low_lim'] == {'feature1': 0.0, 'feature2': 0.0}
assert call_kwargs['upp_lim'] == {'feature1': 100.0, 'feature2': 100.0}