From 44b83943654cd0b74bf7385bb1772b0e90bd7aaf Mon Sep 17 00:00:00 2001
From: Remi Cadene <re.cadene@gmail.com>
Date: Mon, 9 Sep 2024 19:32:35 +0200
Subject: [PATCH] add dynamic import for cv2 and pyrealsense2

---
 .../robot_devices/cameras/intelrealsense.py   |  2 +-
 tests/mock_opencv.py                          | 41 +++++++++++++
 tests/test_cameras.py                         |  5 ++
 tests/utils.py                                | 61 ++++++-------------
 4 files changed, 66 insertions(+), 43 deletions(-)
 create mode 100644 tests/mock_opencv.py

diff --git a/lerobot/common/robot_devices/cameras/intelrealsense.py b/lerobot/common/robot_devices/cameras/intelrealsense.py
index 4806bf785..21ac0b458 100644
--- a/lerobot/common/robot_devices/cameras/intelrealsense.py
+++ b/lerobot/common/robot_devices/cameras/intelrealsense.py
@@ -347,7 +347,7 @@ class IntelRealSenseCamera:
             return color_image
 
     def read_loop(self):
-        while self.stop_event is None or not self.stop_event.is_set():
+        while not self.stop_event.is_set():
             if self.use_depth:
                 self.color_image, self.depth_map = self.read()
             else:
diff --git a/tests/mock_opencv.py b/tests/mock_opencv.py
new file mode 100644
index 000000000..c54ae285d
--- /dev/null
+++ b/tests/mock_opencv.py
@@ -0,0 +1,41 @@
+import cv2
+import numpy as np
+
+
+class MockVideoCapture(cv2.VideoCapture):
+    image = {
+        "480x640": np.random.randint(0, 256, size=(480, 640, 3), dtype=np.uint8),
+        "720x1280": np.random.randint(0, 256, size=(720, 1280, 3), dtype=np.uint8),
+    }
+
+    def __init__(self, *args, **kwargs):
+        self._mock_dict = {
+            cv2.CAP_PROP_FPS: 30,
+            cv2.CAP_PROP_FRAME_WIDTH: 640,
+            cv2.CAP_PROP_FRAME_HEIGHT: 480,
+        }
+
+    def isOpened(self):  # noqa: N802
+        return True
+
+    def set(self, propId: int, value: float) -> bool:  # noqa: N803
+        self._mock_dict[propId] = value
+        return True
+
+    def get(self, propId: int) -> float:  # noqa: N803
+        value = self._mock_dict[propId]
+        if value == 0:
+            if propId == cv2.CAP_PROP_FRAME_HEIGHT:
+                value = 480
+            elif propId == cv2.CAP_PROP_FRAME_WIDTH:
+                value = 640
+        return value
+
+    def read(self):
+        h = self.get(cv2.CAP_PROP_FRAME_HEIGHT)
+        w = self.get(cv2.CAP_PROP_FRAME_WIDTH)
+        ret = True
+        return ret, self.image[f"{h}x{w}"]
+
+    def release(self):
+        pass
diff --git a/tests/test_cameras.py b/tests/test_cameras.py
index 090eb472c..7c9d9e38a 100644
--- a/tests/test_cameras.py
+++ b/tests/test_cameras.py
@@ -5,6 +5,11 @@ Example usage:
 ```bash
 pytest -sx tests/test_cameras.py::test_camera
 ```
+
+Example of running test on mocked_koch:
+```bash
+python -m pytest -sx -k "mocked_koch" tests/test_cameras.py::test_camera
+```
 """
 
 import numpy as np
diff --git a/tests/utils.py b/tests/utils.py
index 9e07219fb..98b5a82a0 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -14,10 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import platform
+import traceback
 from functools import wraps
 
-import cv2
-import numpy as np
 import pytest
 import torch
 
@@ -187,7 +186,24 @@ def require_robot(func):
             kwargs["robot_type"] = robot_type.replace("mocked_", "")
 
             monkeypatch = request.getfixturevalue("monkeypatch")
-            monkeypatch.setattr(cv2, "VideoCapture", MockVideoCapture)
+
+            try:
+                import cv2
+
+                from tests.mock_opencv import MockVideoCapture
+
+                monkeypatch.setattr(cv2, "VideoCapture", MockVideoCapture)
+            except ImportError:
+                traceback.print_exc()
+
+            try:
+                import pyrealsense2 as rs
+
+                # TODO(rcadene):
+                mock_pipeline = None
+                monkeypatch.setattr(rs, "pipeline", mock_pipeline)
+            except ImportError:
+                traceback.print_exc()
 
         # Run test with a real robot. Skip test if robot connection fails.
         else:
@@ -198,42 +214,3 @@ def require_robot(func):
         return func(*args, **kwargs)
 
     return wrapper
-
-
-class MockVideoCapture(cv2.VideoCapture):
-    image = {
-        "480x640": np.random.randint(0, 256, size=(480, 640, 3), dtype=np.uint8),
-        "720x1280": np.random.randint(0, 256, size=(720, 1280, 3), dtype=np.uint8),
-    }
-
-    def __init__(self, *args, **kwargs):
-        self._mock_dict = {
-            cv2.CAP_PROP_FPS: 30,
-            cv2.CAP_PROP_FRAME_WIDTH: 640,
-            cv2.CAP_PROP_FRAME_HEIGHT: 480,
-        }
-
-    def isOpened(self):  # noqa: N802
-        return True
-
-    def set(self, propId: int, value: float) -> bool:  # noqa: N803
-        self._mock_dict[propId] = value
-        return True
-
-    def get(self, propId: int) -> float:  # noqa: N803
-        value = self._mock_dict[propId]
-        if value == 0:
-            if propId == cv2.CAP_PROP_FRAME_HEIGHT:
-                value = 480
-            elif propId == cv2.CAP_PROP_FRAME_WIDTH:
-                value = 640
-        return value
-
-    def read(self):
-        h = self.get(cv2.CAP_PROP_FRAME_HEIGHT)
-        w = self.get(cv2.CAP_PROP_FRAME_WIDTH)
-        ret = True
-        return ret, self.image[f"{h}x{w}"]
-
-    def release(self):
-        pass