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