tests: Fix alignment tool process_video test

deepfakes · Apr 17, 2024 · ec2a95a · ec2a95a
1 parent 957734d
commit ec2a95a
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/tests/tools/alignments/media_test.py b/tests/tools/alignments/media_test.py
@@ -625,12 +625,13 @@ def test_process_video(self, folder: str) -> None:
         folder : str
             Dummy media folder
         """
-        expected = [{"frame_fullname": "images_000001.png",
+        ext = os.path.splitext(folder)[-1]
+        expected = [{"frame_fullname": f"images_000001{ext}",
                      "frame_name": "images_000001",
-                     "frame_extension": ".png"},
-                    {"frame_fullname": "images_000002.png",
+                     "frame_extension": ext},
+                    {"frame_fullname": f"images_000002{ext}",
                      "frame_name": "images_000002",
-                     "frame_extension": ".png"}]
+                     "frame_extension": ext}]
 
         frames = Frames(folder, None)
         returned = list(frames.process_video())