SIENTIAPDE-1252: Implement save_model activity to save trained models to MLflow with comprehensive error handling and add corresponding unit tests.

2025-10-13 10:18:45 -03:00
parent 3fd5ab79fe
commit b3c749872c
2 changed files with 352 additions and 0 deletions
--- a/tests/activities/test_mlflow.py
+++ b/tests/activities/test_mlflow.py
@@ -299,3 +299,251 @@ async def test_update_production_model_error(mlflow):
        )
    else:
        raise AssertionError('No exception raised')
+
+
+@mark.asyncio
+async def test_save_model_success(mlflow):
+    """Test save_model successfully saves model and artifacts to MLflow."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+    from model_manager.utils.models.train_model_result import TrainModelResult
+
+    # Mock train result
+    params = MagicMock(spec=TrainModelParams)
+    params.experiment_name = 'test_experiment'
+
+    train_result = MagicMock(spec=TrainModelResult)
+    train_result.params = params
+    train_result.run_name = None  # Will be set by get_next_run_name
+
+    # Mock repository methods
+    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
+    mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
+    mlflow.model_monitoring_repository.save_run.return_value = None
+
+    input_data = {
+        **metadata,
+        'train_result': train_result,
+    }
+
+    # Call the method
+    response = await mlflow.save_model(input_data)
+
+    # Verify repository methods were called
+    mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('test_experiment')
+    mlflow.model_monitoring_repository.generate_artifacts.assert_called_once_with(train_result)
+    mlflow.model_monitoring_repository.save_run.assert_called_once_with(train_result)
+
+    # Verify response
+    assert response['success'] is True
+    assert response['result'] == train_result
+    assert response['error_message'] is None
+    assert train_result.run_name == 'test_experiment-1'
+
+
+@mark.asyncio
+async def test_save_model_get_next_run_name_error(mlflow):
+    """Test save_model handles error during get_next_run_name."""
+    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.experiment_name = 'test_experiment'
+
+    train_result = MagicMock(spec=TrainModelResult)
+    train_result.params = params
+
+    # Mock error in get_next_run_name
+    mlflow.model_monitoring_repository.get_next_run_name.side_effect = Exception(
+        'MLflow connection error'
+    )
+
+    input_data = {
+        **metadata,
+        'train_result': train_result,
+    }
+
+    # Call the method
+    response = await mlflow.save_model(input_data)
+
+    # Verify error handling
+    assert response['success'] is False
+    assert response['result'] is None
+    assert 'MLflow connection error' in response['error_message']
+
+    # Verify notification was sent
+    mlflow.send_notification.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='SAVE_MODEL_ERROR',
+        message=ANY,
+        block='save_model',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+
+@mark.asyncio
+async def test_save_model_generate_artifacts_error(mlflow):
+    """Test save_model handles error during generate_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.experiment_name = 'test_experiment'
+
+    train_result = MagicMock(spec=TrainModelResult)
+    train_result.params = params
+
+    # Mock successful get_next_run_name but error in generate_artifacts
+    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
+    mlflow.model_monitoring_repository.generate_artifacts.side_effect = FileNotFoundError(
+        'Reports directory does not exist'
+    )
+
+    input_data = {
+        **metadata,
+        'train_result': train_result,
+    }
+
+    # Call the method
+    response = await mlflow.save_model(input_data)
+
+    # Verify error handling
+    assert response['success'] is False
+    assert response['result'] is None
+    assert 'Reports directory does not exist' in response['error_message']
+
+    # Verify notification was sent
+    mlflow.send_notification.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='SAVE_MODEL_ERROR',
+        message=ANY,
+        block='save_model',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+
+@mark.asyncio
+async def test_save_model_save_run_error(mlflow):
+    """Test save_model handles error during save_run."""
+    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.experiment_name = 'test_experiment'
+
+    train_result = MagicMock(spec=TrainModelResult)
+    train_result.params = params
+
+    # Mock successful get_next_run_name and generate_artifacts but error in save_run
+    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
+    mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
+    mlflow.model_monitoring_repository.save_run.side_effect = ValueError(
+        'One or more metrics (MSE, R2, MAE) are None'
+    )
+
+    input_data = {
+        **metadata,
+        'train_result': train_result,
+    }
+
+    # Call the method
+    response = await mlflow.save_model(input_data)
+
+    # Verify error handling
+    assert response['success'] is False
+    assert response['result'] is None
+    assert 'One or more metrics (MSE, R2, MAE) are None' in response['error_message']
+
+    # Verify notification was sent
+    mlflow.send_notification.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='SAVE_MODEL_ERROR',
+        message=ANY,
+        block='save_model',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+
+@mark.asyncio
+async def test_save_model_missing_metadata(mlflow):
+    """Test save_model handles missing metadata 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.experiment_name = 'test_experiment'
+
+    train_result = MagicMock(spec=TrainModelResult)
+    train_result.params = params
+
+    # Mock repository methods
+    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
+    mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
+    mlflow.model_monitoring_repository.save_run.return_value = None
+
+    # Input data without metadata
+    input_data = {
+        'train_result': train_result,
+    }
+
+    # Call the method
+    response = await mlflow.save_model(input_data)
+
+    # Verify it still works (metadata defaults to {})
+    assert response['success'] is True
+    assert response['result'] == train_result
+    assert response['error_message'] is None
+
+
+@mark.asyncio
+async def test_save_model_complete_flow(mlflow):
+    """Test save_model complete flow with all steps."""
+    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.experiment_name = 'production_model'
+
+    train_result = MagicMock(spec=TrainModelResult)
+    train_result.params = params
+    train_result.run_name = None
+    train_result.run_dir = None
+    train_result.report_path = None
+
+    # Mock complete flow
+    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'production_model-5'
+
+    # After generate_artifacts, paths should be set
+    updated_result = MagicMock(spec=TrainModelResult)
+    updated_result.params = params
+    updated_result.run_name = 'production_model-5'
+    updated_result.run_dir = '/reports/production_model-5_20231010'
+    updated_result.report_path = '/reports/production_model-5_20231010/report.html'
+    updated_result.train_data_path = '/reports/production_model-5_20231010/train_data.csv'
+    updated_result.test_data_path = '/reports/production_model-5_20231010/test_data.csv'
+
+    mlflow.model_monitoring_repository.generate_artifacts.return_value = updated_result
+    mlflow.model_monitoring_repository.save_run.return_value = None
+
+    input_data = {
+        **metadata,
+        'train_result': train_result,
+    }
+
+    # Call the method
+    response = await mlflow.save_model(input_data)
+
+    # Verify complete flow
+    mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('production_model')
+    mlflow.model_monitoring_repository.generate_artifacts.assert_called_once()
+    mlflow.model_monitoring_repository.save_run.assert_called_once_with(updated_result)
+
+    # Verify response
+    assert response['success'] is True
+    assert response['result'] == updated_result
+    assert response['error_message'] is None
+    assert updated_result.run_name == 'production_model-5'
+    assert updated_result.run_dir is not None
+    assert updated_result.report_path is not None