From b2ef523c3bd5a6019803ec06aa722ac8ca014ace Mon Sep 17 00:00:00 2001
From: vitor-aignosi <vitor.santos@aignosi.com.br>
Date: Wed, 17 Sep 2025 10:07:32 -0300
Subject: [PATCH] SIENTIAPDE-1222  Update test cases in model repository and
 prediction process

- Replaced string data with MagicMock in test_transform_success and test_transform_error to improve test isolation.
- Updated the predict method calls in test_predict_success and test_predict_error to reflect changes in argument structure.
- Added 'prediction_store_policy' to the test_run configuration in test_prediction_process for consistency with recent updates.
---
 .../utils/repository/test_model_repository.py | 21 ++++++++++++-------
 .../subworkflows/test_prediction_process.py   |  1 +
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py
index 51a3223..5dab80d 100644
--- a/tests/laborious/utils/repository/test_model_repository.py
+++ b/tests/laborious/utils/repository/test_model_repository.py
@@ -113,7 +113,7 @@ def test_detect_and_parse_datetime_index_valid_format(mlflow_repository, data, e
 
 
 def test_transform_success(mlflow_repository):
-    data = 'data'
+    data = MagicMock()
     model_name = 'model'
 
     mlflow_repository.detect_and_parse_datetime_index = MagicMock()
@@ -134,7 +134,7 @@ def test_transform_success(mlflow_repository):
 
 
 def test_transform_error(mlflow_repository):
-    data = 'data'
+    data = MagicMock()
     model_name = 'model'
 
     mlflow_repository.model_serving.get_cached_transform.side_effect = Exception(
@@ -167,10 +167,11 @@ def test_predict_success(mlflow_repository):
         [2, 3]
     )
 
-    output = mlflow_repository.predict(model_name, data, 1)
+    output = mlflow_repository.predict(
+        model_name, data, {}, metadata['metadata'])
 
     mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
-        model_name, data, 1)
+        model_name, data, 0, 'pyfunc', False, 'model')
 
     assert output['success'] is True
     assert output['content'] == {
@@ -185,17 +186,23 @@ def test_predict_success(mlflow_repository):
 
 
 def test_predict_error(mlflow_repository):
-    data = 'data'
+    data = DataFrame({
+        'feat_1': {
+            'index_1': 2,
+            'index_2': 3
+        }
+    })
     model_name = 'model'
 
     mlflow_repository.model_serving.get_cached_predict = MagicMock(
         side_effect=Exception('error')
     )
 
-    output = mlflow_repository.predict(model_name, data, 1)
+    output = mlflow_repository.predict(
+        model_name, data, {}, metadata['metadata'])
 
     mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
-        model_name, data, 1)
+        model_name, data, 0, 'pyfunc', False, 'model')
 
     assert output == {
         'success': False,
diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py
index e76b586..eea8c53 100644
--- a/tests/laborious/workflows/subworkflows/test_prediction_process.py
+++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py
@@ -37,6 +37,7 @@ async def test_run(workflow_mock, prediction_process):
         'model_retention': '30',
         'path_priority': ['continue', 'repeat', 'stop'],
         'opc_output_config': {'test': 'config'},
+        'prediction_store_policy': 'lts:1'
     }
 
     # Mock the activity responses