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.

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", exist_ok=True)
-                makedirs("temp")
 
             file_path = f"temp/raw_data_{model_name}.csv"
             data.to_csv(file_path, index=True)

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"

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
 

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
+        )
+    ])