import tempfile import unittest from pathlib import Path from tests.anyimage.server.support import load_server_module, _make_fake_av class MediaInputTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.media = load_server_module("media.input") def test_collect_images_is_sorted_and_filters_suffixes(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "b.PNG").touch() (root / "a.jpg").touch() (root / "notes.txt").touch() images = self.media.collect_images(root) self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"]) def test_ai_tasks_support_cooperative_cancellation(self): with self.assertRaises(self.media.TaskCancelled): self.media.check_cancelled(lambda: True) self.media.check_cancelled(lambda: False) def test_prepares_a_single_image_file(self): with tempfile.TemporaryDirectory() as directory: image_path = Path(directory) / "input.png" image_path.touch() images, temporary_directory = self.media.prepare_input( image_path, 16, directory, ) self.assertEqual(images, [image_path]) self.assertIsNone(temporary_directory) def test_prepares_a_frame_directory(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "frame-0001.jpg").touch() (root / "frame-0002.png").touch() (root / "notes.txt").touch() images, temporary_directory = self.media.prepare_input( root, 16, directory, ) self.assertEqual( [path.name for path in images], ["frame-0001.jpg", "frame-0002.png"], ) self.assertIsNone(temporary_directory) def test_prepares_first_frames_of_a_frame_directory(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) for index in range(1, 6): (root / f"frame-{index:04d}.jpg").touch() images, temporary_directory = self.media.prepare_input( root, 3, directory, ) self.assertEqual( [path.name for path in images], ["frame-0001.jpg", "frame-0002.jpg", "frame-0003.jpg"], ) self.assertIsNone(temporary_directory) def test_prepares_first_frames_after_offset_in_a_frame_directory(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) for index in range(1, 6): (root / f"frame-{index:04d}.jpg").touch() images, temporary_directory = self.media.prepare_input( root, 3, directory, start_frame=2, ) self.assertEqual( [path.name for path in images], ["frame-0002.jpg", "frame-0003.jpg", "frame-0004.jpg"], ) self.assertIsNone(temporary_directory) def test_extract_video_frames_reads_frames_sequentially(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) fake_av = _make_fake_av(100) image_paths = self.media.extract_video_frames( root / "clip.mkv", root / "out", 10, fake_av, ) self.assertEqual( [path.name for path in image_paths], [f"frame.{index:04d}.jpg" for index in range(1, 11)], ) self.assertEqual(len(fake_av.containers), 1) self.assertEqual(fake_av.containers[0].position, 10) def test_extract_video_frames_starts_after_offset(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) fake_av = _make_fake_av(100) image_paths = self.media.extract_video_frames( root / "clip.mkv", root / "out", 10, fake_av, start_frame=71, ) extracted = [ Path(path).read_text(encoding="utf-8") for path in image_paths ] self.assertEqual(extracted, [str(frame) for frame in range(70, 80)]) self.assertEqual(fake_av.containers[0].position, 80) def test_extract_video_frames_tolerates_short_video_tail(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) image_paths = self.media.extract_video_frames( root / "clip.mkv", root / "out", 10, _make_fake_av(99), ) self.assertEqual( [path.name for path in image_paths], [f"frame.{index:04d}.jpg" for index in range(1, 11)], ) def test_rejects_unsupported_input_file(self): with tempfile.TemporaryDirectory() as directory: input_path = Path(directory) / "input.txt" input_path.touch() with self.assertRaisesRegex(RuntimeError, "Unsupported input format: .txt"): self.media.prepare_input(input_path, 16, directory)