"""Unit tests for Training activities.""" from contextlib import contextmanager from unittest.mock import MagicMock, patch import pandas as pd import pytest from model_manager.utils.models.train_model_params import TrainModelParams from model_manager.utils.models.train_model_result import TrainModelResult def _minimal_params_dict(): return { 'variable_columns': ['a'], 'target_variable': 't', 'bucket_name': 'b', 'file_name': 'f.csv', 'line_separator': '\n', 'decimal_separator': '.', 'date_column': 'timestamp', 'date_format': 'yyyy-MM-dd HH:mm:ss', 'train_size': 80, 'shuffle': True, 'random_state': 42, 'experiment_run_id': 1, 'model_name': 'Linear Regression', 'val_file_name': None, 'data_model_kwargs': {}, 'model_kwargs': {}, 'opt_params': {}, 'model_type': 'linear_regression', 'model_id': None, 'model_metadata': None, } @pytest.fixture def training(): from model_manager.activities.training import Training return Training( mlflow_repository=MagicMock(), plugin_store=MagicMock(), minio_repository=MagicMock(), logger=MagicMock(), notification_handler=MagicMock(), metrics_controller=MagicMock(), ) def test_load_model_metadata_success(training): training.plugin_store.get_model_index = MagicMock( return_value={'schemas': {'components': {'schemas': {}}}} ) inp = {**_minimal_params_dict(), 'metadata': {'w': '1'}} out = training.load_model_metadata(inp) assert 'model_metadata' in out assert out['model_metadata']['schemas'] def test_load_model_metadata_notifies_on_error(training): training.plugin_store.get_model_index = MagicMock(side_effect=RuntimeError('idx')) training.send_notification = MagicMock() inp = {**_minimal_params_dict(), 'metadata': {}} with pytest.raises(RuntimeError, match='idx'): training.load_model_metadata(inp) training.send_notification.assert_called_once() def test_validate_train_params_success(training): pdict = _minimal_params_dict() pdict['model_metadata'] = {'schemas': {'components': {'schemas': {}}}} inp = {**pdict, 'metadata': {}} out = training.validate_train_params(inp) assert isinstance(out, dict) assert out['target_variable'] == 't' def test_validate_train_params_notifies(training): training.send_notification = MagicMock() inp = {'metadata': {}, 'experiment_run_id': 1} with pytest.raises((KeyError, ValueError, TypeError)): training.validate_train_params(inp) training.send_notification.assert_called_once() def test_train_model_download_fails_notifies(training): """train_model notifies and re-raises when MinIO download fails.""" tp = TrainModelParams.from_dict( { **_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}, } ) training.minio_repository.download_file = MagicMock(side_effect=OSError('minio')) training.send_notification = MagicMock() with pytest.raises(OSError, match='minio'): training.train_model({'metadata': {'pod': 'x'}, 'train_params': tp.to_dict()}) training.send_notification.assert_called_once() def test_cleanup_resources(training): training.data_manager_repository.cleanup_run_directory = MagicMock() training.cleanup_resources({'metadata': {}, 'run_dir': '/tmp/x'}) training.data_manager_repository.cleanup_run_directory.assert_called_once_with('/tmp/x', {}) def test_cleanup_resources_notifies_on_error(training): training.data_manager_repository.cleanup_run_directory = MagicMock( side_effect=RuntimeError('rm') ) training.send_notification = MagicMock() with pytest.raises(RuntimeError, match='rm'): training.cleanup_resources({'metadata': {'pod': 'p'}, 'run_dir': '/tmp/x'}) training.send_notification.assert_called_once() @patch('model_manager.activities.training.mlflow') def test_train_model_success_serializes_result(mock_mlflow, training): """Exercise train_model happy path with mocks (MinIO, plugin wrapper, MLflow).""" tp = TrainModelParams.from_dict( { **_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}, } ) train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]}) val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) training.minio_repository.download_file = MagicMock(return_value=b'csv') training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) def _set_metrics(x, _w, **_kw): x.mse_val = 0.1 x.mae_val = 0.2 x.r2_val = 0.9 return x training.data_manager_repository.compute_regression_metrics = MagicMock( side_effect=_set_metrics ) def _fill_report(x, **_kw): x.report_path = '/tmp/report.html' x.train_data_path = '/tmp/train.csv' x.test_data_path = '/tmp/test.csv' x.run_dir = '/tmp/run' return x training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report) wrapper = MagicMock() wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)]) pred_train = pd.DataFrame({'p': [1.0, 2.0]}) pred_val = pd.DataFrame({'p': [1.0]}) wrapper.predict = MagicMock(side_effect=[(pred_train, None), (pred_val, None)]) wrapper.store_model = MagicMock() training.plugin_store.get_model = MagicMock(return_value=wrapper) @contextmanager def _run_ctx(*_a, **_k): info = MagicMock() info.run_name = 'run-n' info.run_id = 'run-i' yield info training.mlflow_repository.start_run = _run_ctx out = training.train_model({'metadata': {'pod': 'p'}, 'train_params': tp.to_dict()}) assert out['run_name'] is None assert out['run_id'] == 'run-i' assert out['run_dir'] == '/tmp/run' mock_mlflow.log_param.assert_any_call('mse_val', 0.1) mock_mlflow.log_param.assert_any_call('mae_val', 0.2) mock_mlflow.log_param.assert_any_call('r2_val', 0.9) mock_mlflow.log_artifact.assert_called() @patch('model_manager.activities.training.mlflow') def test_train_model_without_logger_does_not_set_wrapper_logger(_mock_mlflow, training): """Covers branch where activity logger is None.""" training.logger = None tp = TrainModelParams.from_dict( { **_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}, } ) train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]}) val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) training.minio_repository.download_file = MagicMock(return_value=b'csv') training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) training.data_manager_repository.compute_regression_metrics = MagicMock( side_effect=lambda x, _w, **_kw: x ) def _fill_report(x, **_kw): x.report_path = '/tmp/report.html' x.train_data_path = '/tmp/train.csv' x.test_data_path = '/tmp/test.csv' x.equation_path = '/tmp/eq.json' x.run_dir = '/tmp/run' return x training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report) wrapper = MagicMock() wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)]) wrapper.predict = MagicMock( side_effect=[(pd.DataFrame({'p': [1.0, 2.0]}), None), (pd.DataFrame({'p': [1.0]}), None)] ) wrapper.store_model = MagicMock() training.plugin_store.get_model = MagicMock(return_value=wrapper) @contextmanager def _run_ctx(*_a, **_k): info = MagicMock() info.run_name = 'n' info.run_id = 'i' yield info training.mlflow_repository.start_run = _run_ctx training.train_model({'metadata': {}, 'train_params': tp.to_dict()}) @patch('model_manager.activities.training.mlflow') def test_train_model_train_params_as_dict(mock_mlflow, training): """train_params may arrive as dict and is coerced via TrainModelParams.from_dict.""" d = { **_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}, } train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]}) val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tp = TrainModelParams.from_dict(d) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) training.minio_repository.download_file = MagicMock(return_value=b'csv') training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) training.data_manager_repository.compute_regression_metrics = MagicMock( side_effect=lambda x, _w, **_kw: x ) def _fill_report2(x, **_kw): x.report_path = '/tmp/report.html' x.train_data_path = '/tmp/train.csv' x.test_data_path = '/tmp/test.csv' x.run_dir = '/tmp/run' return x training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report2) wrapper = MagicMock() wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)]) wrapper.predict = MagicMock( side_effect=[(pd.DataFrame({'p': [1.0, 2.0]}), None), (pd.DataFrame({'p': [1.0]}), None)] ) wrapper.store_model = MagicMock() training.plugin_store.get_model = MagicMock(return_value=wrapper) @contextmanager def _run_ctx(*_a, **_k): info = MagicMock() info.run_name = 'n' info.run_id = 'i' yield info training.mlflow_repository.start_run = _run_ctx training.train_model({'metadata': {}, 'train_params': d}) mock_mlflow.log_artifact.assert_called() @patch('model_manager.activities.training.mlflow') def test_train_model_downloads_validation_file_when_set(mock_mlflow, training): """Second MinIO download when val_file_name is set (covers val_bytes branch).""" d = { **_minimal_params_dict(), 'val_file_name': 'val.csv', 'model_metadata': {'schemas': {'components': {'schemas': {}}}}, } tp = TrainModelParams.from_dict(d) train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]}) val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) def _dl(object_name, **_kwargs): if object_name == tp.file_name: return b'train' if object_name == 'val.csv': return b'val' raise AssertionError(object_name) training.minio_repository.download_file = MagicMock(side_effect=_dl) training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) training.data_manager_repository.compute_regression_metrics = MagicMock( side_effect=lambda x, _w, **_kw: x ) def _fill(x, **_kw): x.report_path = '/tmp/report.html' x.train_data_path = '/tmp/train.csv' x.test_data_path = '/tmp/test.csv' x.run_dir = '/tmp/run' return x training.data_manager_repository.generate_report = MagicMock(side_effect=_fill) wrapper = MagicMock() wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)]) wrapper.predict = MagicMock( side_effect=[(pd.DataFrame({'p': [1.0, 2.0]}), None), (pd.DataFrame({'p': [1.0]}), None)] ) wrapper.store_model = MagicMock() training.plugin_store.get_model = MagicMock(return_value=wrapper) @contextmanager def _run_ctx(*_a, **_k): info = MagicMock() info.run_name = 'n' info.run_id = 'i' yield info training.mlflow_repository.start_run = _run_ctx training.train_model({'metadata': {}, 'train_params': tp.to_dict()}) assert training.minio_repository.download_file.call_count == 2 mock_mlflow.log_artifact.assert_called() def test_prepare_data_observes_lag_on_success(training): tp = TrainModelParams.from_dict( {**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}} ) train_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) training.observe_lag_sync = MagicMock() training.emit_metric_sync = MagicMock() from model_manager import metrics as mm_metrics training._prepare_data(b'csv', None, tp, {}) training.observe_lag_sync.assert_called_once() call_args = training.observe_lag_sync.call_args assert call_args.args[1] is mm_metrics.SIENTIA_TRAINING_DATA_PREPARATION_LAG training.emit_metric_sync.assert_not_called() def test_prepare_data_increments_error_counter_and_still_observes_lag_on_failure(training): tp = TrainModelParams.from_dict( {**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}} ) training.data_manager_repository.prepare_training_data = MagicMock( side_effect=RuntimeError('prep-fail') ) training.observe_lag_sync = MagicMock() training.emit_metric_sync = MagicMock() from model_manager import metrics as mm_metrics with pytest.raises(RuntimeError, match='prep-fail'): training._prepare_data(b'csv', None, tp, {}) training.observe_lag_sync.assert_called_once() training.emit_metric_sync.assert_called_once() call_args = training.emit_metric_sync.call_args assert ( call_args.kwargs['metric_object'] is mm_metrics.SIENTIA_TRAINING_DATA_PREPARATION_ERROR_COUNT_TOTAL ) def test_fit_model_observes_lag_on_success(training): tp = TrainModelParams.from_dict( {**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}} ) train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]}) val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) wrapper = MagicMock() wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)]) wrapper.predict = MagicMock( side_effect=[(pd.DataFrame({'p': [1.0, 2.0]}), None), (pd.DataFrame({'p': [1.0]}), None)] ) training.observe_lag_sync = MagicMock() training.emit_metric_sync = MagicMock() from model_manager import metrics as mm_metrics training._fit_model(wrapper, tmr, tp, {}) training.observe_lag_sync.assert_called_once() call_args = training.observe_lag_sync.call_args assert call_args.args[1] is mm_metrics.SIENTIA_TRAINING_MODEL_FIT_LAG training.emit_metric_sync.assert_not_called() def test_fit_model_increments_error_counter_on_failure(training): tp = TrainModelParams.from_dict( {**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}} ) train_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) wrapper = MagicMock() wrapper.train = MagicMock(side_effect=RuntimeError('fit-fail')) training.observe_lag_sync = MagicMock() training.emit_metric_sync = MagicMock() from model_manager import metrics as mm_metrics with pytest.raises(RuntimeError, match='fit-fail'): training._fit_model(wrapper, tmr, tp, {}) training.observe_lag_sync.assert_called_once() training.emit_metric_sync.assert_called_once() call_args = training.emit_metric_sync.call_args assert ( call_args.kwargs['metric_object'] is mm_metrics.SIENTIA_TRAINING_MODEL_FIT_ERROR_COUNT_TOTAL ) @patch('model_manager.activities.training.mm_metrics') @patch('model_manager.activities.training.mlflow') def test_train_model_sets_quality_gauges_after_compute_metrics( mock_mlflow, mock_mm_metrics, training ): tp = TrainModelParams.from_dict( {**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}} ) train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]}) val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) training.minio_repository.download_file = MagicMock(return_value=b'csv') training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) def _set_metrics(x, _w, **_kw): x.mse_val = 0.5 x.mae_val = 0.3 x.r2_val = -0.1 return x training.data_manager_repository.compute_regression_metrics = MagicMock( side_effect=_set_metrics ) def _fill_report(x, **_kw): x.report_path = '/tmp/r.html' x.train_data_path = '/tmp/tr.csv' x.test_data_path = '/tmp/te.csv' x.run_dir = '/tmp/run' return x training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report) wrapper = MagicMock() wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)]) wrapper.predict = MagicMock( side_effect=[(pd.DataFrame({'p': [1.0, 2.0]}), None), (pd.DataFrame({'p': [1.0]}), None)] ) wrapper.store_model = MagicMock() training.plugin_store.get_model = MagicMock(return_value=wrapper) @contextmanager def _run_ctx(*_a, **_k): info = MagicMock() info.run_id = 'rid' yield info training.mlflow_repository.start_run = _run_ctx training.observe_lag_sync = MagicMock() training.emit_metric_sync = MagicMock() training.train_model({'metadata': {}, 'train_params': tp.to_dict()}) mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MSE.labels.return_value.set.assert_called_once_with( 0.5 ) mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MAE.labels.return_value.set.assert_called_once_with( 0.3 ) mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_R2.labels.return_value.set.assert_called_once_with( -0.1 ) @patch('model_manager.activities.training.mm_metrics') @patch('model_manager.activities.training.mlflow') def test_train_model_skips_quality_gauges_when_none(_mock_mlflow, mock_mm_metrics, training): tp = TrainModelParams.from_dict( {**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}} ) train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]}) val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) training.minio_repository.download_file = MagicMock(return_value=b'csv') training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) training.data_manager_repository.compute_regression_metrics = MagicMock( side_effect=lambda x, _w, **_kw: x ) def _fill_report(x, **_kw): x.report_path = '/tmp/r.html' x.train_data_path = '/tmp/tr.csv' x.test_data_path = '/tmp/te.csv' x.run_dir = '/tmp/run' return x training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report) wrapper = MagicMock() wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)]) wrapper.predict = MagicMock( side_effect=[(pd.DataFrame({'p': [1.0, 2.0]}), None), (pd.DataFrame({'p': [1.0]}), None)] ) wrapper.store_model = MagicMock() training.plugin_store.get_model = MagicMock(return_value=wrapper) @contextmanager def _run_ctx(*_a, **_k): info = MagicMock() info.run_id = 'rid' yield info training.mlflow_repository.start_run = _run_ctx training.observe_lag_sync = MagicMock() training.emit_metric_sync = MagicMock() training.train_model({'metadata': {}, 'train_params': tp.to_dict()}) mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MSE.labels.return_value.set.assert_not_called() mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MAE.labels.return_value.set.assert_not_called() mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_R2.labels.return_value.set.assert_not_called() @patch('model_manager.activities.training.mm_metrics') @patch('model_manager.activities.training.mlflow') def test_train_model_increments_trained_total_on_success(_mock_mlflow, mock_mm_metrics, training): tp = TrainModelParams.from_dict( {**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}} ) train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]}) val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) training.minio_repository.download_file = MagicMock(return_value=b'csv') training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) training.data_manager_repository.compute_regression_metrics = MagicMock( side_effect=lambda x, _w, **_kw: x ) def _fill_report(x, **_kw): x.report_path = '/tmp/r.html' x.train_data_path = '/tmp/tr.csv' x.test_data_path = '/tmp/te.csv' x.run_dir = '/tmp/run' return x training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report) wrapper = MagicMock() wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)]) wrapper.predict = MagicMock( side_effect=[(pd.DataFrame({'p': [1.0, 2.0]}), None), (pd.DataFrame({'p': [1.0]}), None)] ) wrapper.store_model = MagicMock() training.plugin_store.get_model = MagicMock(return_value=wrapper) @contextmanager def _run_ctx(*_a, **_k): info = MagicMock() info.run_id = 'rid' yield info training.mlflow_repository.start_run = _run_ctx training.observe_lag_sync = MagicMock() training.emit_metric_sync = MagicMock() training.train_model({'metadata': {}, 'train_params': tp.to_dict()}) training.emit_metric_sync.assert_called_once_with( metric_object=mock_mm_metrics.SIENTIA_TRAINING_MODEL_TRAINED_TOTAL, tags=training._get_training_labels(tp), ) def test_train_model_does_not_increment_trained_total_on_failure(training): tp = TrainModelParams.from_dict( {**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}} ) training.minio_repository.download_file = MagicMock(side_effect=RuntimeError('dl-fail')) training.send_notification = MagicMock() training.emit_metric_sync = MagicMock() with pytest.raises(RuntimeError): training.train_model({'metadata': {}, 'train_params': tp.to_dict()}) training.emit_metric_sync.assert_not_called() def test_train_model_value_error_when_paths_missing_after_report(training): """Raises ValueError when report paths are not populated after generate_report.""" tp = TrainModelParams.from_dict( { **_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}, } ) train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]}) val_df = pd.DataFrame({'a': [1.0], 't': [1.0]}) tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df) training.minio_repository.download_file = MagicMock(return_value=b'x') training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr) training.data_manager_repository.compute_regression_metrics = MagicMock( side_effect=lambda x, _w, **_kw: x ) training.data_manager_repository.generate_report = MagicMock(return_value=tmr) wrapper = MagicMock() wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)]) wrapper.predict = MagicMock( side_effect=[(pd.DataFrame({'p': [1.0, 2.0]}), None), (pd.DataFrame({'p': [1.0]}), None)] ) training.plugin_store.get_model = MagicMock(return_value=wrapper) @contextmanager def _run_ctx(*_a, **_k): info = MagicMock() info.run_name = 'n' info.run_id = 'i' yield info training.mlflow_repository.start_run = _run_ctx training.send_notification = MagicMock() with pytest.raises(ValueError, match='Report path'): training.train_model({'metadata': {}, 'train_params': tp.to_dict()})