diff --git a/tests/utils/repository/test_model_repository.py b/tests/utils/repository/test_model_repository.py index c354262..8d16521 100644 --- a/tests/utils/repository/test_model_repository.py +++ b/tests/utils/repository/test_model_repository.py @@ -199,3 +199,291 @@ def test_get_reports_directory(mock_model_serving_class, mock_logger): assert result.endswith(os.path.join('model_manager', 'reports')) assert os.path.isabs(result) + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +def test_init_artifacts_data_success(mock_model_serving_class, mock_logger, mock_train_result): + """Test _init_artifacts_data successfully prepares data.""" + from model_manager.utils.repository.model_repository import ModelRepository + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + reference_data, current_data = repo._init_artifacts_data(mock_train_result) + + # Check reference data + assert 'target' in reference_data.columns + assert 'prediction' in reference_data.columns + assert len(reference_data) == 3 + + # Check current data + assert 'target' in current_data.columns + assert 'prediction' in current_data.columns + assert len(current_data) == 3 + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +def test_init_artifacts_data_empty_x_train(mock_model_serving_class, mock_logger, mock_train_result): + """Test _init_artifacts_data raises ValueError when x_train is empty.""" + from model_manager.utils.repository.model_repository import ModelRepository + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + mock_train_result.x_train = pd.DataFrame() + + with pytest.raises(ValueError, match='Training features .* are empty'): + repo._init_artifacts_data(mock_train_result) + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +def test_init_artifacts_data_empty_y_train(mock_model_serving_class, mock_logger, mock_train_result): + """Test _init_artifacts_data raises ValueError when y_train is empty.""" + from model_manager.utils.repository.model_repository import ModelRepository + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + mock_train_result.y_train = pd.Series(dtype=float) + + with pytest.raises(ValueError, match='Training target .* is empty'): + repo._init_artifacts_data(mock_train_result) + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +def test_init_artifacts_data_none_y_pred(mock_model_serving_class, mock_logger, mock_train_result): + """Test _init_artifacts_data raises ValueError when y_pred is None.""" + from model_manager.utils.repository.model_repository import ModelRepository + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + mock_train_result.y_pred = None + + with pytest.raises(ValueError, match='Test predictions .* are None'): + repo._init_artifacts_data(mock_train_result) + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +@patch('model_manager.utils.repository.model_repository.datetime') +@patch('model_manager.utils.repository.model_repository.makedirs') +def test_create_run_directory_success( + mock_makedirs, mock_datetime, mock_model_serving_class, mock_logger +): + """Test _create_run_directory creates directory successfully.""" + from model_manager.utils.repository.model_repository import ModelRepository + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + mock_datetime.now.return_value.strftime.return_value = '20240101_120000_123456' + + result = repo._create_run_directory('/tmp/reports', 'test_run') # noqa: S108 + + expected_path = os.path.join('/tmp/reports', 'test_run_20240101_120000_123456') # noqa: S108 + assert result == expected_path + mock_makedirs.assert_called_once_with(expected_path, exist_ok=True) + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +@patch('model_manager.utils.repository.model_repository.makedirs') +def test_create_run_directory_permission_error( + mock_makedirs, mock_model_serving_class, mock_logger +): + """Test _create_run_directory raises PermissionError.""" + from model_manager.utils.repository.model_repository import ModelRepository + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + mock_makedirs.side_effect = PermissionError('Permission denied') + + with pytest.raises(PermissionError, match='Permission denied when creating directory'): + repo._create_run_directory('/tmp/reports', 'test_run') # noqa: S108 + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +@patch('model_manager.utils.repository.model_repository.shutil.copy') +@patch('builtins.open', create=True) +def test_setup_run_directory_success(mock_open, mock_copy, mock_model_serving_class, mock_logger): + """Test _setup_run_directory creates files successfully.""" + from model_manager.utils.repository.model_repository import ModelRepository + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + repo._setup_run_directory('/tmp/test_run', '/tmp/header.html') # noqa: S108 + + # Check that empty files were created + assert mock_open.call_count == 3 + mock_copy.assert_called_once() + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +@patch('model_manager.utils.repository.model_repository.json.dump') +@patch('model_manager.utils.repository.model_repository.Reports') +@patch('builtins.open', create=True) +def test_generate_report_with_equation( + mock_open, mock_reports_class, mock_json_dump, mock_model_serving_class, mock_logger, mock_train_result +): + """Test _generate_report creates equation JSON artifact.""" + from model_manager.utils.repository.model_repository import ModelRepository + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + # Add equation to train result + mock_train_result.equation = { + 'target_variable': 'target', + 'coefficients': {'var1': 1.5, 'var2': -0.75}, + 'intercept': 10.5, + 'equation_string': 'target = 10.5 + 1.5 * var1 + -0.75 * var2', + 'latex_equation': 'target = 10.5 + 1.5 \\cdot var1 + -0.75 \\cdot var2', + 'model_type': 'Linear Regression', + } + + reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]}) + current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]}) + + # Mock DataFrame.to_csv to avoid file I/O + with patch.object(pd.DataFrame, 'to_csv'): + result = repo._generate_report(reference_data, current_data, mock_train_result) + + # Verify equation path was set + assert result.equation_path == os.path.join(mock_train_result.run_dir, 'model_equation.json') + + # Verify JSON was written + mock_json_dump.assert_called() + call_args = mock_json_dump.call_args + assert call_args[0][0] == mock_train_result.equation + assert call_args[1]['indent'] == 2 + assert call_args[1]['ensure_ascii'] is False + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +@patch('model_manager.utils.repository.model_repository.Reports') +@patch('builtins.open', create=True) +def test_generate_report_without_equation( + mock_open, mock_reports_class, mock_model_serving_class, mock_logger, mock_train_result +): + """Test _generate_report works without equation.""" + from model_manager.utils.repository.model_repository import ModelRepository + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + # No equation + mock_train_result.equation = None + + reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]}) + current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]}) + + # Mock DataFrame.to_csv to avoid file I/O + with patch.object(pd.DataFrame, 'to_csv'): + result = repo._generate_report(reference_data, current_data, mock_train_result) + + # Verify equation path was not set + assert not hasattr(result, 'equation_path') or result.equation_path is None + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +@patch('model_manager.utils.repository.model_repository.path.exists') +def test_save_run_with_equation( + mock_exists, mock_model_serving_class, mock_logger, mock_train_result +): + """Test _save_run logs equation artifact.""" + from model_manager.utils.repository.model_repository import ModelRepository + + mock_model_serving_instance = MagicMock() + mock_model_serving_class.return_value = mock_model_serving_instance + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + # Set equation path + mock_train_result.equation_path = '/tmp/test_run/model_equation.json' # noqa: S108 + + # Mock all path.exists calls to return True + mock_exists.return_value = True + + repo._save_run(mock_train_result) + + # Verify equation artifact was logged + logged_artifacts = [ + call[0][0] for call in mock_model_serving_instance.log_artifact.call_args_list + ] + assert '/tmp/test_run/model_equation.json' in logged_artifacts # noqa: S108 + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +@patch('model_manager.utils.repository.model_repository.path.exists') +def test_save_run_without_equation( + mock_exists, mock_model_serving_class, mock_logger, mock_train_result +): + """Test _save_run works without equation.""" + from model_manager.utils.repository.model_repository import ModelRepository + + mock_model_serving_instance = MagicMock() + mock_model_serving_class.return_value = mock_model_serving_instance + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + # No equation + mock_train_result.equation_path = None + + # Mock path.exists to return True for required artifacts + mock_exists.return_value = True + + repo._save_run(mock_train_result) + + # Verify only 3 artifacts were logged (report, train_data, test_data) + assert mock_model_serving_instance.log_artifact.call_count == 3 + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +@patch('model_manager.utils.repository.model_repository.path.exists') +def test_save_run_missing_report( + mock_exists, mock_model_serving_class, mock_logger, mock_train_result +): + """Test _save_run raises ValueError when report is missing.""" + from model_manager.utils.repository.model_repository import ModelRepository + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + # Mock report doesn't exist + def exists_side_effect(path): + return not path.endswith('report.html') + + mock_exists.side_effect = exists_side_effect + + with pytest.raises(ValueError, match='Report file does not exist'): + repo._save_run(mock_train_result) + + +@patch('model_manager.utils.repository.model_repository.ModelServing') +def test_save_run_none_metrics(mock_model_serving_class, mock_logger, mock_train_result): + """Test _save_run raises ValueError when metrics are None.""" + from model_manager.utils.repository.model_repository import ModelRepository + + repo = ModelRepository( + url='http://mlflow.test', username='user', password='pass', logger=mock_logger + ) + + mock_train_result.mse_val = None + + with pytest.raises(ValueError, match='One or more metrics .* are None'): + repo._save_run(mock_train_result) diff --git a/tests/utils/repository/test_training_repository.py b/tests/utils/repository/test_training_repository.py index cea9ace..a67bb7a 100644 --- a/tests/utils/repository/test_training_repository.py +++ b/tests/utils/repository/test_training_repository.py @@ -30,8 +30,6 @@ def sample_params(): return TrainModelParams( experiment_run_id=1, experiment_name='test_experiment', - username='test_user', - model_type='Linear Regression', target_variable='target', variable_columns=['var1', 'var2', 'var3'], lag_train=0, @@ -105,8 +103,6 @@ class TestExtractModelEquation: params = TrainModelParams( experiment_run_id=1, experiment_name='test', - username='test_user', - model_type='Linear Regression', target_variable='y', variable_columns=['x'], lag_train=0, @@ -224,6 +220,78 @@ class TestInitDataPreprocessor: assert preprocessor.lag_transform[col] == 3 +class TestInitScalerDict: + """Tests for _init_scaler_dict method.""" + + def test_init_scaler_dict_without_scaler(self, training_repo, sample_params): + """Test scaler dict initialization when scaler is not used.""" + from model_manager.sientia.models import DataPreprocessor + + sample_params.use_scaler = False + process_data = MagicMock(spec=DataPreprocessor) + + result = training_repo._init_scaler_dict(process_data, sample_params) + + assert result == {} + + def test_init_scaler_dict_with_minmax_scaler(self, training_repo, sample_params): + """Test scaler dict initialization with MinMaxScaler.""" + from model_manager.sientia.models import DataPreprocessor + from sientia_do.operations.normalization import MinMaxScaler + + sample_params.use_scaler = True + + # Create mock MinMaxScaler + mock_scaler = MagicMock(spec=MinMaxScaler) + mock_scaler.x_min = np.array([0.0, 1.0, 2.0]) + mock_scaler.x_max = np.array([10.0, 11.0, 12.0]) + mock_scaler.y_min = 0.5 + mock_scaler.y_max = 100.5 + + # Create mock preprocessor that returns the scaler + process_data = MagicMock(spec=DataPreprocessor) + process_data.get_scaler.return_value = mock_scaler + + result = training_repo._init_scaler_dict(process_data, sample_params) + + # Check feature scalers + assert 'var1' in result + assert 'var2' in result + assert 'var3' in result + assert result['var1'] == {'min': 0.0, 'max': 10.0} + assert result['var2'] == {'min': 1.0, 'max': 11.0} + assert result['var3'] == {'min': 2.0, 'max': 12.0} + + # Check target scaler + assert 'target' in result + assert result['target'] == {'min': 0.5, 'max': 100.5} + + def test_init_scaler_dict_with_z_scaler(self, training_repo, sample_params): + """Test scaler dict initialization with Z_Scaler.""" + from model_manager.sientia.models import DataPreprocessor + from sientia_do.operations.normalization import Z_Scaler + + sample_params.use_scaler = True + + # Create mock Z_Scaler + mock_scaler = MagicMock(spec=Z_Scaler) + expected_dict = { + 'var1': {'mean': 5.0, 'std': 1.5}, + 'var2': {'mean': 10.0, 'std': 2.0}, + 'target': {'mean': 50.0, 'std': 10.0}, + } + mock_scaler.create_dict.return_value = expected_dict + + # Create mock preprocessor + process_data = MagicMock(spec=DataPreprocessor) + process_data.get_scaler.return_value = mock_scaler + + result = training_repo._init_scaler_dict(process_data, sample_params) + + assert result == expected_dict + mock_scaler.create_dict.assert_called_once() + + class TestAfterTrainCalculation: """Tests for after_train_calculation method.""" @@ -307,3 +375,93 @@ class TestAfterTrainCalculation: for call in mock_logger.info.call_args_list ) + def test_after_train_with_custom_scaler_denormalization( + self, training_repo, sample_params, sample_linear_model + ): + """Test denormalization with custom scaler that has denormalize methods.""" + sample_params.use_scaler = True + + # Create mock data + x_train = pd.DataFrame( + {'var1': [1, 2, 3], 'var2': [4, 5, 6], 'var3': [7, 8, 9]}, index=[0, 1, 2] + ) + x_test = pd.DataFrame({'var1': [10], 'var2': [11], 'var3': [12]}, index=[3]) + y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target') + y_test = pd.Series([400], index=[3], name='target') + + # Mock predict + sample_linear_model.predict = MagicMock(return_value=np.array([450.0])) + + # Create mock scaler with denormalize methods + mock_scaler = MagicMock() + mock_scaler.denormalize_single_input = MagicMock(side_effect=lambda x, col: x * 2) + mock_scaler.denormalize_predictions = MagicMock(return_value=np.array([900.0])) + + # Create mock preprocessor + mock_process_data = MagicMock() + mock_process_data.get_scaler.return_value = mock_scaler + + train_result = TrainModelResult( + params=sample_params, + process_data=mock_process_data, + x_train=x_train, + x_test=x_test, + y_train=y_train, + y_test=y_test, + regr=sample_linear_model, + scaler_dict={}, + ) + + result = training_repo.after_train_calculation(sample_params, train_result) + + # Verify denormalize methods were called + assert mock_scaler.denormalize_single_input.called + assert mock_scaler.denormalize_predictions.called + assert result.y_pred is not None + + def test_after_train_with_sklearn_scaler( + self, training_repo, sample_params, sample_linear_model + ): + """Test denormalization with sklearn StandardScaler.""" + sample_params.use_scaler = True + + # Create mock data + x_train = pd.DataFrame( + {'var1': [1, 2, 3], 'var2': [4, 5, 6], 'var3': [7, 8, 9]}, index=[0, 1, 2] + ) + x_test = pd.DataFrame({'var1': [10], 'var2': [11], 'var3': [12]}, index=[3]) + y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target') + y_test = pd.Series([400], index=[3], name='target') + + # Mock predict + sample_linear_model.predict = MagicMock(return_value=np.array([450.0])) + + # Create mock sklearn scaler (without denormalize methods) + mock_scaler = MagicMock() + # Remove denormalize methods to trigger sklearn path + if hasattr(mock_scaler, 'denormalize_single_input'): + delattr(mock_scaler, 'denormalize_single_input') + mock_scaler.inverse_transform = MagicMock(side_effect=lambda x: x * 2) + + # Create mock preprocessor with feature_names_order + mock_process_data = MagicMock() + mock_process_data.get_scaler.return_value = mock_scaler + mock_process_data.feature_names_order = ['var1', 'var2', 'var3'] + + train_result = TrainModelResult( + params=sample_params, + process_data=mock_process_data, + x_train=x_train, + x_test=x_test, + y_train=y_train, + y_test=y_test, + regr=sample_linear_model, + scaler_dict={}, + ) + + result = training_repo.after_train_calculation(sample_params, train_result) + + # Verify inverse_transform was called + assert mock_scaler.inverse_transform.called + assert result.y_pred is not None +