Add test for loading DLC 3D pose data

Author: lochhh

tests/fixtures/datasets.py

@@ -312,6 +312,24 @@ def valid_dlc_poses_df():
     return pd.read_hdf(pytest.DATA_PATHS.get("DLC_single-wasp.predictions.h5"))
 
 
+@pytest.fixture
+def valid_dlc_3d_poses_df(valid_dlc_poses_df):
+    """Mock and return a valid DLC-style 3D poses DataFrame.
+
+    The only difference between 2D and 3D DLC DataFrames is that
+    the coordinate level in the columns MultiIndex includes 'z' instead of
+    'likelihood'.
+    """
+    cols = [
+        (scorer, bodypart, "z" if coord == "likelihood" else coord)
+        for scorer, bodypart, coord in valid_dlc_poses_df.columns.to_list()
+    ]
+    valid_dlc_poses_df.columns = pd.MultiIndex.from_tuples(
+        cols, names=valid_dlc_poses_df.columns.names
+    )
+    return valid_dlc_poses_df
+
+
 # -------------------- Invalid bboxes datasets --------------------
 @pytest.fixture
 def missing_var_bboxes_dataset(valid_bboxes_dataset):

tests/test_unit/test_io/test_load_poses.py

@@ -99,15 +99,20 @@ def test_load_from_dlc_file(file_name, helpers):
 
 
 @pytest.mark.parametrize(
-    "source_software", ["DeepLabCut", "LightningPose", None]
+    "poses_df_fixture, source_software",
+    [
+        ("valid_dlc_poses_df", "DeepLabCut"),
+        ("valid_dlc_poses_df", "LightningPose"),
+        ("valid_dlc_poses_df", None),
+        ("valid_dlc_3d_poses_df", "DeepLabCut"),
+    ],
 )
-def test_load_from_dlc_style_df(valid_dlc_poses_df, source_software, helpers):
-    """Test that loading pose tracks from a valid DLC-style DataFrame
-    returns a proper Dataset.
-    """
-    ds = load_poses.from_dlc_style_df(
-        valid_dlc_poses_df, source_software=source_software
-    )
+def test_load_from_dlc_style_df(
+    poses_df_fixture, source_software, helpers, request
+):
+    """Test loading pose tracks from DLC-style DataFrames (2D and 3D)."""
+    df = request.getfixturevalue(poses_df_fixture)
+    ds = load_poses.from_dlc_style_df(df, source_software=source_software)
     expected_values = {
         **expected_values_poses,
         "source_software": source_software,