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.

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"