diff --git a/model_manager/utils/filters/mlflow_filters.py b/model_manager/utils/filters/mlflow_filters.py
index 82970e8..493ed9c 100644
--- a/model_manager/utils/filters/mlflow_filters.py
+++ b/model_manager/utils/filters/mlflow_filters.py
@@ -54,8 +54,8 @@ def nan_values_filter(predictions: DataFrame, _config: dict) -> bool:
     """
     data = (
         predictions.replace({None: np.nan})
+        .infer_objects(copy=False)
         .drop(columns=['timestamp'], errors='ignore')
-        .infer_objects()
     )
 
     if data.isna().all().all():
diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py
index 58fcf7e..116f1e4 100644
--- a/tests/laborious/utils/repository/test_model_repository.py
+++ b/tests/laborious/utils/repository/test_model_repository.py
@@ -194,7 +194,7 @@ def test_get_experiment_success(mlflow, mlflow_repository):
 
     output = mlflow_repository.get_experiment('test')
 
-    assert output == 0
+    assert output == '0'
 
 
 @patch('model_manager.utils.repository.model_repository.mlflow')
@@ -245,7 +245,7 @@ def test_get_experiment_last_run_error(mlflow, mlflow_repository):
 @patch('model_manager.utils.repository.model_repository.mlflow.sklearn')
 @patch('model_manager.utils.repository.model_repository.mlflow.set_experiment')
 def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
-    mlflow_repository.model_serving.get_model_run_id = MagicMock(return_value='0')
+    mlflow_repository.model_serving.get_model_info = MagicMock(return_value='0')
     mlflow_repository.model_serving.get_model_uri = MagicMock(return_value='test')
     mlflow_repository.get_experiment_by_run_id = MagicMock()
 
@@ -268,9 +268,7 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
 
     output = mlflow_repository.create_model_experiment('test', data)
 
-    mlflow_repository.model_serving.get_model_run_id.assert_called_once_with(
-        'test', stage='Production'
-    )
+    mlflow_repository.model_serving.get_model_info.assert_called_once_with('test')
     mlflow_repository.model_serving.get_model_uri.assert_called_once_with('0', prediction=False)
 
     sklearn.load_model.assert_has_calls(
diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
index fb83fac..e1fd744 100644
--- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
+++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
@@ -69,7 +69,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
                 {
                     'schema': input_data['schema'],
                     'table_name': input_data['table_name'],
-                    'data': workflow_mock.execute_activity_method.return_value,
+                    'data': workflow_mock.execute_local_activity_method.return_value,
                     **metadata,
                     'timestamp_conversion': {
                         'column': 'timestamp',
@@ -130,7 +130,7 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
                 {
                     'schema': input_data['schema'],
                     'table_name': input_data['table_name'],
-                    'data': workflow_mock.execute_activity_method.return_value,
+                    'data': workflow_mock.execute_local_activity_method.return_value,
                     **metadata,
                     'timestamp_conversion': {
                         'column': 'timestamp',
@@ -143,5 +143,5 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
         ]
     )
 
-    assert workflow_mock.execute_activity_method.call_count == 3
+    assert workflow_mock.execute_activity_method.call_count == 2
     assert workflow_mock.execute_local_activity_method.call_count == 1