From cd4d2257d31d4302e62de442e72fda28d0ac024e Mon Sep 17 00:00:00 2001
From: Remi Cadene <remi.cadene@huggingface.co>
Date: Thu, 12 Sep 2024 02:43:21 +0200
Subject: [PATCH] Fix unit test

---
 tests/mock_intelrealsense.py | 18 ++++++++++++++++++
 tests/test_cameras.py        |  6 ++++--
 tests/test_control_robot.py  | 11 +++++------
 tests/utils.py               |  3 ++-
 4 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/tests/mock_intelrealsense.py b/tests/mock_intelrealsense.py
index 1e6c3f5f..c1d6add2 100644
--- a/tests/mock_intelrealsense.py
+++ b/tests/mock_intelrealsense.py
@@ -108,3 +108,21 @@ class MockDepthFrame:
 
     def get_data(self):
         return np.ones((self.config.height, self.config.width), dtype=np.uint16)
+
+
+class MockDevice:
+    def __init__(self):
+        pass
+
+    def get_info(self, camera_info) -> str:
+        del camera_info  # unused
+        # return fake serial number
+        return "123456789"
+
+
+class MockContext:
+    def __init__(self):
+        pass
+
+    def query_devices(self):
+        return [MockDevice()]
diff --git a/tests/test_cameras.py b/tests/test_cameras.py
index 161d9a19..0c6ad08a 100644
--- a/tests/test_cameras.py
+++ b/tests/test_cameras.py
@@ -8,14 +8,16 @@ Example of running a specific test:
 pytest -sx tests/test_cameras.py::test_camera
 ```
 
-Example of running test on a real OpenCV camera connected to the computer:
+Example of running test on a real camera connected to the computer:
 ```bash
 pytest -sx tests/test_cameras.py::test_camera[opencv]
+pytest -sx tests/test_cameras.py::test_camera[intelrealsense]
 ```
 
-Example of running test on a mocked version of an OpenCV camera:
+Example of running test on a mocked version of the camera:
 ```bash
 pytest -sx -k "mocked_opencv" tests/test_cameras.py::test_camera
+pytest -sx -k "mocked_intelrealsense" tests/test_cameras.py::test_camera
 ```
 """
 
diff --git a/tests/test_control_robot.py b/tests/test_control_robot.py
index 406edeb4..cd70263b 100644
--- a/tests/test_control_robot.py
+++ b/tests/test_control_robot.py
@@ -2,15 +2,14 @@ from pathlib import Path
 
 import pytest
 
-from lerobot import available_robots
 from lerobot.common.policies.factory import make_policy
 from lerobot.common.utils.utils import init_hydra_config
 from lerobot.scripts.control_robot import calibrate, record, replay, teleoperate
 from tests.test_robots import make_robot
-from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, require_robot
+from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, TEST_MOTOR_TYPES, require_robot
 
 
-@pytest.mark.parametrize("robot_type", available_robots)
+@pytest.mark.parametrize("robot_type", TEST_MOTOR_TYPES)
 @require_robot
 def test_teleoperate(request, robot_type):
     robot = make_robot(robot_type)
@@ -20,7 +19,7 @@ def test_teleoperate(request, robot_type):
     del robot
 
 
-@pytest.mark.parametrize("robot_type", available_robots)
+@pytest.mark.parametrize("robot_type", TEST_MOTOR_TYPES)
 @require_robot
 def test_calibrate(request, robot_type):
     robot = make_robot(robot_type)
@@ -28,7 +27,7 @@ def test_calibrate(request, robot_type):
     del robot
 
 
-@pytest.mark.parametrize("robot_type", available_robots)
+@pytest.mark.parametrize("robot_type", TEST_MOTOR_TYPES)
 @require_robot
 def test_record_without_cameras(tmpdir, request, robot_type):
     root = Path(tmpdir)
@@ -38,7 +37,7 @@ def test_record_without_cameras(tmpdir, request, robot_type):
     record(robot, fps=30, root=root, repo_id=repo_id, warmup_time_s=1, episode_time_s=1, num_episodes=2)
 
 
-@pytest.mark.parametrize("robot_type", available_robots)
+@pytest.mark.parametrize("robot_type", TEST_MOTOR_TYPES)
 @require_robot
 def test_record_and_replay_and_policy(tmpdir, request, robot_type):
     env_name = "koch_real"
diff --git a/tests/utils.py b/tests/utils.py
index c7915ea2..b1eb052e 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -298,12 +298,13 @@ def mock_cameras(request):
     try:
         import pyrealsense2 as rs
 
-        from tests.mock_intelrealsense import MockConfig, MockFormat, MockPipeline, MockStream
+        from tests.mock_intelrealsense import MockConfig, MockContext, MockFormat, MockPipeline, MockStream
 
         monkeypatch.setattr(rs, "config", MockConfig)
         monkeypatch.setattr(rs, "pipeline", MockPipeline)
         monkeypatch.setattr(rs, "stream", MockStream)
         monkeypatch.setattr(rs, "format", MockFormat)
+        monkeypatch.setattr(rs, "context", MockContext)
     except ImportError:
         traceback.print_exc()