small fix unit tests

This commit is contained in:
Remi Cadene 2024-10-14 17:08:08 +02:00
parent 87fef10f89
commit e0394961cb
1 changed files with 3 additions and 3 deletions

View File

@ -319,7 +319,7 @@ def test_record_with_event_rerecord_episode(tmpdir, request, robot_type, mock):
robot = make_robot(robot_type, overrides=overrides, mock=mock) robot = make_robot(robot_type, overrides=overrides, mock=mock)
with ( with (
patch("lerobot.scripts.control_robot.init_keyboard_listener") as mock_listener, patch("lerobot.scripts.control_robot.init_keyboard_listener") as mock_listener,
patch("lerobot.common.robot_devices.control_robot.add_frame", wraps=add_frame) as mock_add_frame, patch("lerobot.common.robot_devices.control_utils.add_frame", wraps=add_frame) as mock_add_frame,
): ):
mock_events = {} mock_events = {}
mock_events["exit_early"] = True mock_events["exit_early"] = True
@ -368,7 +368,7 @@ def test_record_with_event_exit_early(tmpdir, request, robot_type, mock):
robot = make_robot(robot_type, overrides=overrides, mock=mock) robot = make_robot(robot_type, overrides=overrides, mock=mock)
with ( with (
patch("lerobot.scripts.control_robot.init_keyboard_listener") as mock_listener, patch("lerobot.scripts.control_robot.init_keyboard_listener") as mock_listener,
patch("lerobot.common.robot_devices.control_robot.add_frame", wraps=add_frame) as mock_add_frame, patch("lerobot.common.robot_devices.control_utils.add_frame", wraps=add_frame) as mock_add_frame,
): ):
mock_events = {} mock_events = {}
mock_events["exit_early"] = True mock_events["exit_early"] = True
@ -418,7 +418,7 @@ def test_record_with_event_stop_recording(tmpdir, request, robot_type, mock, num
robot = make_robot(robot_type, overrides=overrides, mock=mock) robot = make_robot(robot_type, overrides=overrides, mock=mock)
with ( with (
patch("lerobot.scripts.control_robot.init_keyboard_listener") as mock_listener, patch("lerobot.scripts.control_robot.init_keyboard_listener") as mock_listener,
patch("lerobot.common.robot_devices.control_robot.add_frame", wraps=add_frame) as mock_add_frame, patch("lerobot.common.robot_devices.control_utils.add_frame", wraps=add_frame) as mock_add_frame,
): ):
mock_events = {} mock_events = {}
mock_events["exit_early"] = True mock_events["exit_early"] = True