test(cleaning): cleaning up tests

tests/datasets/test_aggregate.py

@@ -33,6 +33,7 @@ from tests.fixtures.constants import (
     DUMMY_CAMERA_FEATURES,
     DUMMY_DEPTH_CAMERA_FEATURES,
     DUMMY_REPO_ID,
+    DUMMY_CAMERA_FEATURES_WITH_DEPTH,
 )
 
 CAMERA_FEATURES_WITH_DEPTH = {**DUMMY_CAMERA_FEATURES, **DUMMY_DEPTH_CAMERA_FEATURES}
@@ -304,14 +305,14 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
         repo_id=f"{DUMMY_REPO_ID}_0",
         total_episodes=ds_0_num_episodes,
         total_frames=ds_0_num_frames,
-        camera_features=CAMERA_FEATURES_WITH_DEPTH,
+        camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
     )
     ds_1 = lerobot_dataset_factory(
         root=tmp_path / "test_1",
         repo_id=f"{DUMMY_REPO_ID}_1",
         total_episodes=ds_1_num_episodes,
         total_frames=ds_1_num_frames,
-        camera_features=CAMERA_FEATURES_WITH_DEPTH,
+        camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
     )
 
     # Confirm depth was actually wired into the source datasets so the
@@ -476,14 +477,14 @@ def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory):
         repo_id=f"{DUMMY_REPO_ID}_small_0",
         total_episodes=ds_0_num_episodes,
         total_frames=ds_0_num_frames,
-        camera_features=CAMERA_FEATURES_WITH_DEPTH,
+        camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
     )
     ds_1 = lerobot_dataset_factory(
         root=tmp_path / "small_1",
         repo_id=f"{DUMMY_REPO_ID}_small_1",
         total_episodes=ds_1_num_episodes,
         total_frames=ds_1_num_frames,
-        camera_features=CAMERA_FEATURES_WITH_DEPTH,
+        camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
     )
 
     assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key"
@@ -549,7 +550,7 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory):
             repo_id=f"{DUMMY_REPO_ID}_regression_{i}",
             total_episodes=2,
             total_frames=100,
-            camera_features=CAMERA_FEATURES_WITH_DEPTH,
+            camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
         )
         datasets.append(ds)
 
@@ -711,7 +712,7 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
         total_episodes=ds_0_num_episodes,
         total_frames=ds_0_num_frames,
         use_videos=False,
-        camera_features=CAMERA_FEATURES_WITH_DEPTH,
+        camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
     )
     ds_1 = lerobot_dataset_factory(
         root=tmp_path / "image_1",
@@ -719,7 +720,7 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
         total_episodes=ds_1_num_episodes,
         total_frames=ds_1_num_frames,
         use_videos=False,
-        camera_features=CAMERA_FEATURES_WITH_DEPTH,
+        camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
     )
 
     # Verify source datasets have image keys

tests/datasets/test_dataset_tools.py

@@ -37,7 +37,8 @@ from lerobot.datasets.dataset_tools import (
     split_dataset,
 )
 from lerobot.datasets.io_utils import load_info
-from tests.datasets.test_video_encoding import _add_frames, require_h264, require_libsvtav1
+from tests.datasets.test_video_encoding import require_h264, require_libsvtav1
+from tests.fixtures.dataset_factories import add_frames
 
 
 @pytest.fixture
@@ -1350,9 +1351,9 @@ def test_reencode_dataset_multi_key_multiprocessing(
         camera_encoder=initial_cfg,
     )
 
-    _add_frames(dataset, num_frames=4)
+    add_frames(dataset, num_frames=4)
     dataset.save_episode()
-    _add_frames(dataset, num_frames=4)
+    add_frames(dataset, num_frames=4)
     dataset.save_episode()
     dataset.finalize()
 

tests/fixtures/constants.py

@@ -57,5 +57,6 @@ DUMMY_DEPTH_CAMERA_FEATURES = {
         "info": DUMMY_DEPTH_VIDEO_INFO,
     },
 }
+DUMMY_CAMERA_FEATURES_WITH_DEPTH = {**DUMMY_CAMERA_FEATURES, **DUMMY_DEPTH_CAMERA_FEATURES}
 DUMMY_CHW = (3, 96, 128)
 DUMMY_HWC = (96, 128, 3)

tests/fixtures/dataset_factories.py

@@ -38,8 +38,6 @@ from lerobot.datasets.utils import (
     DEFAULT_VIDEO_PATH,
     DatasetInfo,
 )
-from lerobot.datasets.video_utils import encode_video_frames
-from lerobot.utils.constants import DEFAULT_FEATURES
 from tests.fixtures.constants import (
     DEFAULT_FPS,
     DUMMY_CAMERA_FEATURES,
@@ -47,6 +45,44 @@ from tests.fixtures.constants import (
     DUMMY_REPO_ID,
     DUMMY_ROBOT_TYPE,
 )
+from lerobot.datasets.video_utils import encode_video_frames
+from lerobot.utils.constants import DEFAULT_FEATURES
+
+
+def add_frames(
+    dataset: LeRobotDataset, num_frames: int
+) -> None:
+    """Append ``num_frames`` synthetic frames to ``dataset``.
+
+    Generates per-feature payloads from ``dataset.meta``: uint16 depth ramps for
+    keys in ``dataset.meta.depth_keys``, uint8 random noise for video/image keys,
+    and float32 zeros for everything else. ``DEFAULT_FEATURES`` (timestamp,
+    frame_index, ...) are auto-populated by ``add_frame`` and skipped here.
+    """
+    if video_keys is None:
+        video_keys = dataset.meta.video_keys
+    depth_keys = set(dataset.meta.depth_keys)
+    # Smooth gradient base reused per (H, W) to keep depth frames cheap to
+    # encode (HEVC Main 12 hates white noise).
+    _depth_base_cache: dict[tuple[int, int], np.ndarray] = {}
+    for i in range(num_frames):
+        frame: dict = {"task": "test"}
+        for key, ft in dataset.meta.features.items():
+            if key in DEFAULT_FEATURES:
+                continue
+            shape = ft["shape"]
+            if key in depth_keys:
+                h, w, _ = shape
+                base = _depth_base_cache.setdefault(
+                    (h, w),
+                    np.linspace(100.0, 10_000.0, h * w, dtype=np.float32).reshape(h, w, 1),
+                )
+                frame[key] = (base + 50.0 * i).clip(0, 65535).astype(np.uint16)
+            elif key in video_keys:
+                frame[key] = np.random.randint(0, 256, shape, dtype=np.uint8)
+            else:
+                frame[key] = np.zeros(shape, dtype=np.float32)
+        dataset.add_frame(frame)
 
 
 class LeRobotDatasetFactory(Protocol):