diff --git a/cli/medperf/tests/commands/dataset/test_create.py b/cli/medperf/tests/commands/dataset/test_create.py
index 82eb52ed9..49e82a355 100644
--- a/cli/medperf/tests/commands/dataset/test_create.py
+++ b/cli/medperf/tests/commands/dataset/test_create.py
@@ -195,15 +195,30 @@ def test_fails_if_invalid_params(self, mocker, benchmark_uid, cube_uid, comms, u
         else:
             preparation.validate()
 
-    @pytest.mark.parametrize("in_path", ["data_path", "input_path", "/usr/data/path"])
-    @pytest.mark.parametrize("out_path", ["out_path", "~/.medperf/data/123"])
+    @pytest.mark.parametrize(
+        "in_path",
+        [
+            ["data", "labels"],
+            ["in_data", "in_labels"],
+            ["/usr/data/path", "usr/labels/path"],
+        ],
+    )
+    @pytest.mark.parametrize(
+        "out_path",
+        [
+            ["out_data", "out_labels"],
+            ["~/.medperf/data/123/data", "~/.medperf/data/123/labels"],
+        ],
+    )
     def test_generate_uids_assigns_uids_to_obj_properties(
         self, mocker, in_path, out_path, preparation
     ):
         # Arrange
         mocker.patch(PATCH_DATAPREP.format("get_folders_hash"), side_effect=lambda x: x)
-        preparation.data_path = in_path
-        preparation.out_datapath = out_path
+        preparation.data_path = in_path[0]
+        preparation.labels_path = in_path[1]
+        preparation.out_datapath = out_path[0]
+        preparation.out_labelspath = out_path[1]
 
         # Act
         preparation.generate_uids()
diff --git a/cli/medperf/tests/commands/result/test_create.py b/cli/medperf/tests/commands/result/test_create.py
index f19647baa..8dfe90606 100644
--- a/cli/medperf/tests/commands/result/test_create.py
+++ b/cli/medperf/tests/commands/result/test_create.py
@@ -323,6 +323,6 @@ def test_execution_of_one_model_writes_result(self, mocker, setup):
 
         # Assert
         assert (
-            yaml.load(open(expected_file))["results"]
+            yaml.safe_load(open(expected_file))["results"]
             == self.state_variables["models_props"][model_uid]["results"]
         )