tests(all shapes): enhancing tests to cover all possible features shapes

tests/datasets/test_dataset_tools.py

@@ -304,40 +304,45 @@ def test_merge_empty_list(tmp_path):
         merge_datasets([], output_repo_id="merged", output_dir=tmp_path)
 
 
-def test_add_features_with_values(sample_dataset, tmp_path):
+@pytest.mark.parametrize(
-    """Test adding a feature with pre-computed values."""
+    "values, feature_shape, expected_item_shape",
+    [
+        (np.random.randn(50).astype(np.float32), (1,), ()),
+        (np.random.randn(50, 1).astype(np.float32), (1,), ()),
+        (np.random.randn(50, 7).astype(np.float32), (7,), (7,)),
+        (np.random.randn(50, 1, 4).astype(np.float32), (1, 4), (1, 4)),
+    ],
+    ids=["scalar_1d", "scalar_2d", "vector", "matrix"],
+)
+def test_add_features_with_values(sample_dataset, tmp_path, values, feature_shape, expected_item_shape):
+    """Test adding a pre-computed feature across supported per-frame shapes."""
     num_frames = sample_dataset.meta.total_frames
-    reward_values = np.random.randn(num_frames, 1).astype(np.float32)
+    assert len(values) == num_frames
 
-    feature_info = {
+    feature_info = {"dtype": "float32", "shape": feature_shape, "names": None}
-        "dtype": "float32",
+    features = {"new_feature": (values, feature_info)}
-        "shape": (1,),
-        "names": None,
-    }
-    features = {
-        "reward": (reward_values, feature_info),
-    }
 
     with (
         patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
         patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
     ):
         mock_get_safe_version.return_value = "v3.0"
-        mock_snapshot_download.return_value = str(tmp_path / "with_reward")
+        mock_snapshot_download.return_value = str(tmp_path / "with_feature")
 
         new_dataset = add_features(
             dataset=sample_dataset,
             features=features,
-            output_dir=tmp_path / "with_reward",
+            output_dir=tmp_path / "with_feature",
         )
 
-    assert "reward" in new_dataset.meta.features
+    assert "new_feature" in new_dataset.meta.features
-    assert new_dataset.meta.features["reward"] == feature_info
+    assert new_dataset.meta.features["new_feature"] == feature_info
 
     assert len(new_dataset) == num_frames
     sample_item = new_dataset[0]
-    assert "reward" in sample_item
+    assert "new_feature" in sample_item
-    assert isinstance(sample_item["reward"], torch.Tensor)
+    assert isinstance(sample_item["new_feature"], torch.Tensor)
+    assert tuple(sample_item["new_feature"].shape) == expected_item_shape
 
 
 def test_add_features_with_callable(sample_dataset, tmp_path):
@@ -377,41 +382,6 @@ def test_add_features_with_callable(sample_dataset, tmp_path):
     assert first_frame["frame_index"] == 0
     assert float(first_frame["reward"]) == 0.0
 
-
-def test_add_features_with_multidimensional_values(sample_dataset, tmp_path):
-    """Test adding a multi-dimensional per-frame feature."""
-    num_frames = sample_dataset.meta.total_frames
-    feature_values = np.random.randn(num_frames, 1, 4).astype(np.float32)
-
-    feature_info = {
-        "dtype": "float32",
-        "shape": (1, 4),
-        "names": None,
-    }
-    features = {
-        "multi_dim_feature": (feature_values, feature_info),
-    }
-
-    with (
-        patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
-        patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
-    ):
-        mock_get_safe_version.return_value = "v3.0"
-        mock_snapshot_download.return_value = str(tmp_path / "with_multi_dim_feature")
-
-        new_dataset = add_features(
-            dataset=sample_dataset,
-            features=features,
-            output_dir=tmp_path / "with_multi_dim_feature",
-        )
-
-    assert "multi_dim_feature" in new_dataset.meta.features
-    sample_item = new_dataset[0]
-    assert "multi_dim_feature" in sample_item
-    assert isinstance(sample_item["multi_dim_feature"], torch.Tensor)
-    assert tuple(sample_item["multi_dim_feature"].shape) == (1, 4)
-
-
 def test_add_existing_feature(sample_dataset, tmp_path):
     """Test error when adding an existing feature."""
     feature_info = {"dtype": "float32", "shape": (1,)}