From c738e8b0df47294ae676df71beed2c9207c1bc74 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 24 Jul 2025 10:24:08 -0300 Subject: [PATCH] SIENTIAPDE-1171 Refactor model_repository and enhance test coverage for MLFlow functionalities - Updated model_repository to ensure the 'temp' directory is created if it doesn't exist using `exist_ok=True`. - Added new tests for retraining and updating production models, including error handling scenarios. - Improved existing tests for model management workflows to ensure robustness and reliability. --- .../utils/repository/model_repository.py | 3 +- tests/laborious/activities/test_mlflow.py | 133 +++++++++++++++++- .../utils/repository/test_model_repository.py | 59 ++++++-- .../workflows/test_minimal_retrain.py | 97 +++++++++++++ 4 files changed, 280 insertions(+), 12 deletions(-) create mode 100644 tests/laborious/workflows/test_minimal_retrain.py diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 187e324..2f10a9d 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -168,8 +168,7 @@ class MLFlowRepository(): # dynamic parameters, including model itself mlflow.sklearn.log_model(data_model, "data_model") - if not path.exists("temp"): - makedirs("temp") + makedirs("temp", exist_ok=True) file_path = f"temp/raw_data_{model_name}.csv" data.to_csv(file_path, index=True) diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index a92f3f3..3d83736 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -1,6 +1,7 @@ -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import numpy as np +from pandas import DataFrame from pytest import fixture, mark from laborious.activities.mlflow import MLFlow @@ -29,7 +30,7 @@ def test___init__(mock_mlflow_repository): @fixture @patch("laborious.activities.mlflow.MLFlowRepository") def mlflow(mock_mlflow_repository): - return MLFlow( + mlflow = MLFlow( mlflow_host="http://localhost:5000", mlflow_port=5000, mlflow_username="admin", @@ -38,6 +39,10 @@ def mlflow(mock_mlflow_repository): notification_handler=MagicMock() ) + mlflow.send_notification = MagicMock() + + return mlflow + metadata = { "metadata": { @@ -141,3 +146,127 @@ async def test_request_predict(mock_max, mock_dataframe, mlflow): mlflow.model_monitoring_repository.predict.assert_called_once_with( 'test_model', mock_dataframe.return_value, 30 ) + + +@mark.asyncio +async def test_retrain_model(mlflow): + data = { + "model_id": [4, 5, 6, 7], + "created_at": [1, 2, 3, 4], + "timestamp": [1, 1, 2, 2], + "variable": ["var1", "var2", "var1", "var2"], + "value": [1, 2, 3, 4] + } + + mlflow.model_monitoring_repository.retrain_model.return_value = ( + 'Model retrained successfully', 'test') + + response = await mlflow.retrain_model({ + **metadata, + 'data': data, + 'model_name': 'test_model' + }) + + mlflow.model_monitoring_repository.retrain_model.assert_called_once() + + assert response == { + "status": 'Model retrained successfully', + "timestamp": 2, + "experiment": 'test' + } + + +@mark.asyncio +async def test_retrain_model_error(mlflow): + mlflow.model_monitoring_repository.retrain_model.side_effect = Exception( + 'Error retraining model' + ) + + data = { + "model_id": [4, 5, 6, 7], + "created_at": [1, 2, 3, 4], + "timestamp": [1, 1, 2, 2], + "variable": ["var1", "var2", "var1", "var2"], + "value": [1, 2, 3, 4] + } + + try: + await mlflow.retrain_model({ + **metadata, + 'data': data, + 'model_name': 'test_model' + }) + except Exception as e: + assert str(e) == 'Error retraining model' + mlflow.send_notification.assert_called_once_with( + metadata=metadata['metadata'], + notification_id='RETRAIN_MODEL_ERROR', + message='Error retraining model test_model: Error retraining model', + block='retrain_model', + attachment_content=ANY + ) + else: + assert False, "No exception raised" + + +@mark.asyncio +async def test_update_production_model(mlflow): + mlflow.model_monitoring_repository.update_production_model.return_value = ( + { + "data1": 1, + "data2": 2 + } + ) + + input_data = { + **metadata, + 'model_name': 'test_model', + 'model_id': 1, + 'experiment': 'test', + 'timestamp': 2, + 'status': 'success' + } + + response = await mlflow.update_production_model(input_data) + + mlflow.model_monitoring_repository.update_production_model.assert_called_once_with( + experiment='test', model_name='test_model') + + assert response == { + 'data1': {0: 1}, + 'data2': {0: 2}, + 'model_id': {0: 1}, + 'model_name': {0: 'test_model'}, + 'timestamp': {0: 2}, + 'status': {0: 'success'} + } + + +@mark.asyncio +async def test_update_production_model_error(mlflow): + mlflow.model_monitoring_repository.update_production_model.side_effect = Exception( + 'Error updating production model' + ) + + input_data = { + **metadata, + 'model_name': 'test_model', + 'model_id': 1, + 'experiment': 'test', + 'timestamp': 2, + 'status': 'success' + } + + try: + await mlflow.update_production_model(input_data) + except Exception as e: + assert str(e) == 'Error updating production model' + mlflow.send_notification.assert_called_once_with( + metadata=metadata['metadata'], + notification_id='UPDATE_PRODUCTION_MODEL_ERROR', + message='Error updating production model test_model: Error updating production model', + block='update_production_model', + attachment_content=ANY + ) + else: + assert False, "No exception raised" diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index f3978cb..62f4fc3 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -164,6 +164,18 @@ def test_get_experiment_last_run(mlflow, mlflow_repository): assert output == '2' +@patch('laborious.utils.repository.model_repository.mlflow') +def test_get_experiment_last_run_error(mlflow, mlflow_repository): + mlflow.search_runs.return_value = [] + + try: + mlflow_repository.get_experiment_last_run(0) + except ValueError as e: + assert str(e) == 'Runs is not a pandas DataFrame' + else: + assert False + + @patch('laborious.utils.repository.model_repository.mlflow.sklearn') @patch('laborious.utils.repository.model_repository.mlflow.set_experiment') def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): @@ -233,14 +245,6 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): @patch('laborious.utils.repository.model_repository.mlflow.log_artifact') def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository): - _os_remove = patch( - 'laborious.utils.repository.model_repository.remove') - _os_makedirs = patch( - 'laborious.utils.repository.model_repository.makedirs') - _os_path_exists = patch( - 'laborious.utils.repository.model_repository.path.exists', - MagicMock(return_value=False)) - prediction_model_mock = MagicMock() data_model_mock = MagicMock() experiment = 'test' @@ -277,6 +281,27 @@ def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, ml assert output == ("Model retrained successfully", experiment) +def test_retrain_model(mlflow_repository): + data = MagicMock() + model_name = 'test' + + mlflow_repository.create_model_experiment = MagicMock( + return_value=('data_model', 'prediction_model', '0')) + + mlflow_repository.perform_model_retrain = MagicMock( + return_value='Model retrained successfully') + + output = mlflow_repository.retrain_model(data, model_name) + + mlflow_repository.create_model_experiment.assert_called_once_with( + model_name, data) + + mlflow_repository.perform_model_retrain.assert_called_once_with( + 'data_model', 'prediction_model', '0', model_name, data) + + assert output == 'Model retrained successfully' + + @patch('laborious.utils.repository.model_repository.mlflow') def test_update_production_model_by_run_id(mlflow, mlflow_repository): client_mock = MagicMock() @@ -312,6 +337,24 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository): } +@patch('laborious.utils.repository.model_repository.mlflow') +def test_update_production_model_by_run_id_error(mlflow, mlflow_repository): + mlflow.tracking.MlflowClient.return_value = MagicMock( + get_registered_model=MagicMock( + return_value=MagicMock( + latest_versions={} + ) + ) + ) + + try: + mlflow_repository.update_production_model_by_run_id('0', 'test') + except Exception as e: + assert str(e) == 'Model versions is not a list' + else: + assert False + + def test_update_production_model(mlflow_repository): connector = mlflow_repository diff --git a/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py new file mode 100644 index 0000000..bcb2288 --- /dev/null +++ b/tests/laborious/workflows/test_minimal_retrain.py @@ -0,0 +1,97 @@ +from unittest.mock import AsyncMock, MagicMock, call, patch, ANY +from pytest import fixture, mark +from laborious.activities.activities import Activities +from laborious.workflows.minimal_retrain import MinimalRetrain + + +@fixture +def minimal_retrain() -> MinimalRetrain: + return MinimalRetrain() + + +metadata = { + "metadata": { + "model_id": "test_model_id", + "model_name": "test_model", + "workflow_name": "minimal_retrain", + "schedule_name": "test_schedule", + }, +} + + +@mark.asyncio +@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock) +async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): + input_data = { + "model_id": "test_model_id", + "model_name": "test_model", + "workflow_name": "minimal_retrain", + "schedule_name": "test_schedule", + "query": "test_query", + "schema": "test_schema", + "table_name": "test_table", + } + + workflow_mock.execute_activity_method = AsyncMock( + return_value={ + "data1": "1", + "data2": "2", + } + ) + + await minimal_retrain.run(input_data) + + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.load_custom_query, + { + **metadata, + "query": input_data["query"], + }, + retry_policy=ANY, + start_to_close_timeout=ANY + ) + ] + ) + + workflow_mock.execute_activity_method.assert_has_calls([ + call( + Activities.retrain_model, + { + **metadata, + 'data': workflow_mock.execute_local_activity_method.return_value, + 'model_name': input_data['model_name'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY + ) + ]) + + workflow_mock.execute_activity_method.assert_has_calls([ + call( + Activities.update_production_model, + { + **metadata, + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + **workflow_mock.execute_activity_method.return_value, + }, + retry_policy=ANY, + start_to_close_timeout=ANY + ) + ]) + + workflow_mock.execute_activity_method.assert_has_calls([ + call( + Activities.export_data_to_postgres, + { + **metadata, + 'data': workflow_mock.execute_activity_method.return_value, + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY + ) + ])