small fix unit tests

huggingface · Oct 14, 2024 · cde4945 · cde4945
1 parent e039496
commit cde4945
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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"