add unit test for conert_video

AllenNeuralDynamics · Nov 6, 2024 · b70512f · b70512f
1 parent 36e1c61
commit b70512f
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/src/aind_behavior_video_transformation/__init__.py b/src/aind_behavior_video_transformation/__init__.py
@@ -9,4 +9,5 @@
 from aind_behavior_video_transformation.transform_videos import (  # noqa F401
     CompressionEnum,
     CompressionRequest,
+    convert_video
 )
diff --git a/tests/test_transform_videos.py b/tests/test_transform_videos.py
@@ -15,6 +15,7 @@
     BehaviorVideoJobSettings,
     CompressionEnum,
     CompressionRequest,
+    convert_video
 )
 
 
@@ -52,6 +53,28 @@ class TestBehaviorVideoJob(unittest.TestCase):
     test_vid_name = "clip.mp4"
     test_vid_path = test_data_path / test_vid_name
 
+    @patch("aind_behavior_video_transformation.etl.time")
+    def test_convert_video(self, mock_time: MagicMock):
+        """Unit test convert video."""
+
+        # Equivalent to CompressionEnum.GAMMA_ENCODING
+        arg_set = ("", "-vf "
+        '"scale=out_color_matrix=bt709:out_range=full:sws_dither=none,'
+        "format=yuv420p10le,colorspace=ispace=bt709:all=bt709:dither=none,"
+        'scale=out_range=tv:sws_dither=none,format=yuv420p" -c:v libx264 '
+        "-preset veryslow -crf 18 -pix_fmt yuv420p "
+        '-metadata author="Allen Institute for Neural Dyamics" '
+        "-movflags +faststart+write_colr")
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_path = Path(temp_dir)
+            compressed_out_path = convert_video(self.test_vid_path,
+                          temp_path,
+                          arg_set)
+
+            out_path = temp_path / self.test_vid_name
+
+            self.assertTrue(str(out_path) == str(compressed_out_path))
+
     @patch("aind_behavior_video_transformation.etl.time")
     def test_run_job(self, mock_time: MagicMock):
         """Tests run_job method."""