diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py
index 42b4d10..a65c6d4 100644
--- a/tests/laborious/activities/test_gates.py
+++ b/tests/laborious/activities/test_gates.py
@@ -379,8 +379,14 @@ async def test_format_prediction_no_timestamp(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'data': {'prediction': [1], 'response_time': [0.1]},
-        'timestamp': '2023-05-26 11:12:27',
+        'data': {
+            'prediction': {
+                '2023-05-26 11:12:27': 1
+            },
+            'response_time': {
+                '2023-05-26 11:12:27': 0.1
+            }
+        },
         'model_id': 'test_model',
         'prediction_confidence': 0.9,
         'prediction_store_policy': 'lts:1'
@@ -404,9 +410,18 @@ async def test_format_prediction_with_timestamp_erl(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'data': {'prediction': [1, 2, 3],
-                 'response_time': [0.1, 0.2, 0.3],
-                 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']},
+        'data': {
+            'prediction': {
+                '2023-05-26 11:12:27': 1,
+                '2023-05-26 11:12:28': 2,
+                '2023-05-26 11:12:29': 3,
+            },
+            'response_time': {
+                '2023-05-26 11:12:27': 0.1,
+                '2023-05-26 11:12:28': 0.2,
+                '2023-05-26 11:12:29': 0.3,
+            }
+        },
         'model_id': 'test_model',
         'prediction_confidence': 0.9,
         'prediction_store_policy': 'erl:2'
@@ -431,9 +446,18 @@ async def test_format_prediction_with_timestamp_lts(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'data': {'prediction': [1, 2, 3],
-                 'response_time': [0.1, 0.2, 0.3],
-                 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']},
+        'data': {
+            'prediction': {
+                '2023-05-26 11:12:27': 1,
+                '2023-05-26 11:12:28': 2,
+                '2023-05-26 11:12:29': 3,
+            },
+            'response_time': {
+                '2023-05-26 11:12:27': 0.1,
+                '2023-05-26 11:12:28': 0.2,
+                '2023-05-26 11:12:29': 0.3,
+            }
+        },
         'model_id': 'test_model',
         'prediction_confidence': 0.9,
         'prediction_store_policy': 'lts:2'
diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py
index 29cbe5c..8b3716e 100644
--- a/tests/laborious/utils/repository/test_model_repository.py
+++ b/tests/laborious/utils/repository/test_model_repository.py
@@ -59,7 +59,12 @@ def test_transform_error(mlflow_repository):
 
 
 def test_predict_success(mlflow_repository):
-    data = 'data'
+    data = DataFrame({
+        'feat_1': {
+            'index_1': 2,
+            'index_2': 3
+        }
+    })
     model_name = 'model'
     mlflow_repository.model_serving.get_cached_predict.return_value = np.array(
         [2, 3]
@@ -71,10 +76,15 @@ def test_predict_success(mlflow_repository):
         model_name, data, 1)
 
     assert output['success'] is True
-    assert output['content'] == {'prediction': {
-        0: 2,
-        1: 3
-    }, 'response_time': ANY}
+    assert output['content'] == {
+        'prediction': {
+            'index_1': 2,
+            'index_2': 3
+        }, 'response_time': {
+            'index_1': ANY,
+            'index_2': ANY
+        }
+    }
 
 
 def test_predict_error(mlflow_repository):