Update test_commands.py

add ExportAnalysisFile `all_frames` arg tests
talmolab · Dec 6, 2023 · 9b69740 · 9b69740
1 parent 750ac13
commit 9b69740
Showing 1 changed file with 42 additions and 3 deletions.
diff --git a/tests/gui/test_commands.py b/tests/gui/test_commands.py
@@ -226,7 +226,7 @@ def assert_videos_written(num_videos: int, labels_path: str = None):
         assert_videos_written(num_videos=1, labels_path=context.state["filename"])
 
         return
-
+    
     # Test with all_videos False (single video)
     params = {"all_videos": False, "csv": csv}
     okay = ExportAnalysisFile_ask(context=context, params=params)
@@ -272,13 +272,52 @@ def assert_videos_written(num_videos: int, labels_path: str = None):
     ExportAnalysisFile.do_action(context=context, params=params)
     assert_videos_written(num_videos=1, labels_path=context.state["filename"])
 
-    # Test with all videos True
+    # Test with all_videos True
     params = {"all_videos": True, "csv": csv}
     okay = ExportAnalysisFile_ask(context=context, params=params)
     assert okay == True
     ExportAnalysisFile.do_action(context=context, params=params)
     assert_videos_written(num_videos=2, labels_path=context.state["filename"])
 
+    # Test with all_videos True and all_frames True
+    params = {"all_videos": True, "all_frames": True, "csv": csv}
+    okay = ExportAnalysisFile_ask(context=context, params=params)
+    assert okay == True
+    ExportAnalysisFile.do_action(context=context, params=params)
+    assert_videos_written(num_videos=2, labels_path=context.state["filename"])
+
+    # Test with all_videos False and all_frames True
+    params = {"all_videos": False, "all_frames": True, "csv": csv}
+    okay = ExportAnalysisFile_ask(context=context, params=params)
+    assert okay == True
+    ExportAnalysisFile.do_action(context=context, params=params)
+    assert_videos_written(num_videos=2, labels_path=context.state["filename"])
+
+    # Test with all_videos False and all_frames False
+    params = {"all_videos": False, "all_frames":False, "csv": csv}
+    okay = ExportAnalysisFile_ask(context=context, params=params)
+    assert okay == True
+    ExportAnalysisFile.do_action(context=context, params=params)
+    assert_videos_written(num_videos=2, labels_path=context.state["filename"])
+
+    # Add labels path and test with all_videos True and all_frames True (single video)
+    context.state["filename"] = str(tmpdir.with_name("path.to.labels"))
+    params = {"all_videos": True, "all_frames": True, "csv": csv}
+    okay = ExportAnalysisFile_ask(context=context, params=params)
+    assert okay == True
+    ExportAnalysisFile.do_action(context=context, params=params)
+    assert_videos_written(num_videos=1, labels_path=context.state["filename"])
+
+    # Add a video (no labels) and test with all_videos True and all_frames True
+    labels.add_video(small_robot_mp4_vid)
+
+    params = {"all_videos": True, "all_frames: True, "csv": csv}
+    okay = ExportAnalysisFile_ask(context=context, params=params)
+    assert okay == True
+    ExportAnalysisFile.do_action(context=context, params=params)
+    assert_videos_written(num_videos=1, labels_path=context.state["filename"])
+
+
     # Test with videos with the same filename
     (tmpdir / "session1").mkdir()
     (tmpdir / "session2").mkdir()
@@ -302,7 +341,7 @@ def assert_videos_written(num_videos: int, labels_path: str = None):
     for video in all_videos:
         labels.remove_video(labels.videos[-1])
 
-    params = {"all_videos": True, "csv": csv}
+    params = {"all_videos": True, "all_frames": True, "csv": csv}
     with pytest.raises(ValueError):
         okay = ExportAnalysisFile_ask(context=context, params=params)