AnyImage/tests/test_media_input.py

import tempfile
import unittest
from pathlib import Path
from server_test_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_copy_color_images_numbers_a_still_image(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            source = root / "source.png"
            source.write_bytes(b"frame")
            output = root / "result"
            output.mkdir()

            self.media.copy_color_images([source], output)

            self.assertEqual(
                (output / "image.0001.png").read_bytes(),
                b"frame",
            )


    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 type"):
                self.media.prepare_input(input_path, 16, directory)