From cde4945b2985787a033b9f7f7a62d18c4943e229 Mon Sep 17 00:00:00 2001 From: Remi Cadene Date: Mon, 14 Oct 2024 17:19:25 +0200 Subject: [PATCH] small fix unit tests --- lerobot/common/robot_devices/control_utils.py | 3 ++- lerobot/scripts/control_robot.py | 3 +++ tests/test_control_robot.py | 2 +- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/lerobot/common/robot_devices/control_utils.py b/lerobot/common/robot_devices/control_utils.py index a93b79498..826ca5de1 100644 --- a/lerobot/common/robot_devices/control_utils.py +++ b/lerobot/common/robot_devices/control_utils.py @@ -316,7 +316,8 @@ def stop_recording(robot, listener, display_cameras): robot.disconnect() if not is_headless(): - listener.stop() + if listener is not None: + listener.stop() if display_cameras: cv2.destroyAllWindows() diff --git a/lerobot/scripts/control_robot.py b/lerobot/scripts/control_robot.py index 9b76ccd66..59c16eef2 100644 --- a/lerobot/scripts/control_robot.py +++ b/lerobot/scripts/control_robot.py @@ -289,6 +289,9 @@ def record( # Increment by one dataset["current_episode_index"] save_current_episode(dataset) + if events["stop_recording"]: + break + log_say("Stop recording", play_sounds, blocking=True) stop_recording(robot, listener, display_cameras) diff --git a/tests/test_control_robot.py b/tests/test_control_robot.py index 4df858b08..2c0bca9b6 100644 --- a/tests/test_control_robot.py +++ b/tests/test_control_robot.py @@ -446,5 +446,5 @@ def test_record_with_event_stop_recording(tmpdir, request, robot_type, mock, num ) assert not mock_events["exit_early"], "`exit_early` wasn't properly reset to False" - assert mock_add_frame.call_count == 1, "`add_frame` should have been called 2 times" + assert mock_add_frame.call_count == 1, "`add_frame` should have been called 1 time" assert len(dataset) == 1, "`dataset` should contain only 1 frame"