From 3cb85bcd4bace7402b7905fb6f3055be1b32b0f7 Mon Sep 17 00:00:00 2001
From: Remi Cadene <remi.cadene@huggingface.co>
Date: Thu, 26 Sep 2024 13:09:08 +0200
Subject: [PATCH] Fix unit test

---
 tests/test_control_robot.py |  1 +
 tests/test_motors.py        | 24 +++++++++++++++++-------
 tests/utils.py              | 10 ++++++----
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/tests/test_control_robot.py b/tests/test_control_robot.py
index 2b35e798..efb32419 100644
--- a/tests/test_control_robot.py
+++ b/tests/test_control_robot.py
@@ -92,6 +92,7 @@ def test_record_and_replay_and_policy(tmpdir, request, robot_type, mock):
         episode_time_s=1,
         num_episodes=2,
         push_to_hub=False,
+        # TODO(rcadene, aliberts): test video=True
         video=False,
     )
 
diff --git a/tests/test_motors.py b/tests/test_motors.py
index cb416989..10380330 100644
--- a/tests/test_motors.py
+++ b/tests/test_motors.py
@@ -30,10 +30,9 @@ import time
 import numpy as np
 import pytest
 
-from lerobot import available_motors
 from lerobot.common.robot_devices.motors.utils import MotorsBus
 from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
-from tests.utils import TEST_MOTOR_TYPES, require_motor
+from tests.utils import TEST_MOTOR_TYPES, mock_input, require_motor
 
 DYNAMIXEL_PORT = "/dev/tty.usbmodem575E0032081"
 DYNAMIXEL_MOTORS = {
@@ -58,19 +57,30 @@ def make_motors_bus(motor_type: str, **kwargs) -> MotorsBus:
         raise ValueError(f"The motor type '{motor_type}' is not valid.")
 
 
-# TODO(rcadene): implement mocked version of this test
-@pytest.mark.parametrize("motor_type, mock", [(m, False) for m in available_motors])
+@pytest.mark.parametrize("motor_type, mock", TEST_MOTOR_TYPES)
 @require_motor
 def test_find_port(request, motor_type, mock):
     from lerobot.common.robot_devices.motors.dynamixel import find_port
 
-    find_port()
+    if mock:
+        # To run find_port without user input
+        monkeypatch = request.getfixturevalue("monkeypatch")
+        monkeypatch.setattr("builtins.input", mock_input)
+
+        with pytest.raises(OSError):
+            find_port()
+    else:
+        find_port()
 
 
-# TODO(rcadene): implement mocked version of this test
-@pytest.mark.parametrize("motor_type, mock", [(m, False) for m in available_motors])
+@pytest.mark.parametrize("motor_type, mock", TEST_MOTOR_TYPES)
 @require_motor
 def test_configure_motors_all_ids_1(request, motor_type, mock):
+    if mock:
+        # To run find_port without user input
+        monkeypatch = request.getfixturevalue("monkeypatch")
+        monkeypatch.setattr("builtins.input", mock_input)
+
     input("Are you sure you want to re-configure the motors? Press enter to continue...")
     # This test expect the configuration was already correct.
     motors_bus = make_motors_bus(motor_type)
diff --git a/tests/utils.py b/tests/utils.py
index aeafd9d1..dd2c904e 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -213,11 +213,8 @@ def require_robot(func):
             mock_cameras(request, camera_type)
             mock_motors(request)
 
-            def mock_input(text):
-                print(text)
-
-            monkeypatch = request.getfixturevalue("monkeypatch")
             # To run calibration without user input
+            monkeypatch = request.getfixturevalue("monkeypatch")
             monkeypatch.setattr("builtins.input", mock_input)
 
         # Run test with a real robot. Skip test if robot connection fails.
@@ -301,6 +298,11 @@ def require_motor(func):
     return wrapper
 
 
+def mock_input(text=None):
+    if text is not None:
+        print(text)
+
+
 def mock_cameras(request, camera_type="all"):
     # TODO(rcadene): Redesign the mocking tests
     monkeypatch = request.getfixturevalue("monkeypatch")