ffx/tests/unit/test_media_descriptor_import_subtitles.py

from __future__ import annotations

from pathlib import Path
import sys
import tempfile
import unittest


SRC_ROOT = Path(__file__).resolve().parents[2] / "src"

if str(SRC_ROOT) not in sys.path:
    sys.path.insert(0, str(SRC_ROOT))


from ffx.logging_utils import get_ffx_logger  # noqa: E402
from ffx.media_descriptor import MediaDescriptor  # noqa: E402
from ffx.track_descriptor import TrackDescriptor  # noqa: E402
from ffx.track_disposition import TrackDisposition  # noqa: E402
from ffx.track_type import TrackType  # noqa: E402


class MediaDescriptorImportSubtitlesTests(unittest.TestCase):
    COMPLETE_SIDECAR_NAMES = (
        "A2_t01_2_deu_DEF",
        "A2_t01_3_eng",
        "A2_t01_4_eng",
    )

    def write_complete_sidecar_set(self, directory: str, extension: str) -> None:
        for basename in self.COMPLETE_SIDECAR_NAMES:
            (Path(directory) / f"{basename}.{extension}").write_text(
                "WEBVTT\n\n",
                encoding="utf-8",
            )

    def make_descriptor(self, indices=(3,)) -> MediaDescriptor:
        return MediaDescriptor(
            context={"logger": get_ffx_logger()},
            track_descriptors=[
                TrackDescriptor(
                    index=index,
                    source_index=index,
                    sub_index=subIndex,
                    track_type=TrackType.SUBTITLE,
                    tags={"language": "eng", "title": "DB Subtitle"},
                    disposition_set={TrackDisposition.DEFAULT},
                )
                for subIndex, index in enumerate(indices)
            ],
        )

    def test_import_subtitles_preserves_target_dispositions_when_requested(self):
        descriptor = self.make_descriptor()

        with tempfile.TemporaryDirectory() as tmpdir:
            sidecar_path = Path(tmpdir) / "dball_S01E01_3_deu_FOR.vtt"
            sidecar_path.write_text("WEBVTT\n\n", encoding="utf-8")

            descriptor.importSubtitles(
                tmpdir,
                "dball",
                season=1,
                episode=1,
                preserve_dispositions=True,
            )

        track = descriptor.getSubtitleTracks()[0]
        self.assertEqual(str(sidecar_path), track.getExternalSourceFilePath())
        self.assertEqual("deu", track.getTags()["language"])
        self.assertEqual({TrackDisposition.DEFAULT}, track.getDispositionSet())

    def test_import_subtitles_uses_sidecar_dispositions_by_default(self):
        descriptor = self.make_descriptor()

        with tempfile.TemporaryDirectory() as tmpdir:
            sidecar_path = Path(tmpdir) / "dball_S01E01_3_deu_FOR.vtt"
            sidecar_path.write_text("WEBVTT\n\n", encoding="utf-8")

            descriptor.importSubtitles(
                tmpdir,
                "dball",
                season=1,
                episode=1,
            )

        track = descriptor.getSubtitleTracks()[0]
        self.assertEqual(str(sidecar_path), track.getExternalSourceFilePath())
        self.assertEqual("deu", track.getTags()["language"])
        self.assertEqual({TrackDisposition.FORCED}, track.getDispositionSet())

    def test_strict_basename_import_recognizes_vtt_asset_set(self):
        descriptor = self.make_descriptor(indices=(2, 3, 4))

        with tempfile.TemporaryDirectory() as tmpdir:
            self.write_complete_sidecar_set(tmpdir, "vtt")
            result = descriptor.importSubtitles(
                tmpdir,
                "A2_t01",
                strict=True,
            )

        self.assertEqual(3, result["candidate_count"])
        self.assertEqual([2, 3, 4], result["imported_track_indices"])
        self.assertEqual([], result["missing_track_indices"])
        self.assertEqual(
            [
                "A2_t01_2_deu_DEF.vtt",
                "A2_t01_3_eng.vtt",
                "A2_t01_4_eng.vtt",
            ],
            [
                Path(track.getExternalSourceFilePath()).name
                for track in descriptor.getSubtitleTracks()
            ],
        )

    def test_strict_basename_import_accepts_dotted_mkv_extension(self):
        descriptor = self.make_descriptor(indices=(2, 3, 4))

        with tempfile.TemporaryDirectory() as tmpdir:
            self.write_complete_sidecar_set(tmpdir, "mkv")
            result = descriptor.importSubtitles(
                tmpdir,
                "A2_t01",
                extension=".mkv",
                strict=True,
            )

        self.assertEqual(3, result["candidate_count"])
        self.assertEqual([2, 3, 4], result["imported_track_indices"])
        self.assertEqual([], result["missing_track_indices"])
        self.assertTrue(
            all(
                track.getExternalSourceFilePath().endswith(".mkv")
                for track in descriptor.getSubtitleTracks()
            )
        )

    def test_strict_basename_import_reports_missing_tracks(self):
        descriptor = self.make_descriptor(indices=(2, 3, 4))

        with tempfile.TemporaryDirectory() as tmpdir:
            sidecarPath = Path(tmpdir) / "episode_2_deu.vtt"
            sidecarPath.write_text("WEBVTT\n\n", encoding="utf-8")

            result = descriptor.importSubtitles(
                tmpdir,
                "episode",
                strict=True,
            )

        self.assertEqual([2], result["imported_track_indices"])
        self.assertEqual([3, 4], result["missing_track_indices"])

    def test_strict_basename_import_rejects_too_many_files(self):
        descriptor = self.make_descriptor(indices=(2,))

        with tempfile.TemporaryDirectory() as tmpdir:
            for filename in ("episode_2_deu.vtt", "episode_3_eng.vtt"):
                (Path(tmpdir) / filename).write_text("WEBVTT\n\n", encoding="utf-8")

            with self.assertRaisesRegex(ValueError, "2 matching .* for 1 subtitle tracks"):
                descriptor.importSubtitles(tmpdir, "episode", strict=True)

    def test_strict_basename_import_rejects_unknown_track_index(self):
        descriptor = self.make_descriptor(indices=(2, 3, 4))

        with tempfile.TemporaryDirectory() as tmpdir:
            (Path(tmpdir) / "episode_9_eng.vtt").write_text(
                "WEBVTT\n\n",
                encoding="utf-8",
            )

            with self.assertRaisesRegex(ValueError, "track index pattern does not match"):
                descriptor.importSubtitles(tmpdir, "episode", strict=True)

    def test_strict_basename_import_rejects_malformed_filtered_filename(self):
        descriptor = self.make_descriptor(indices=(2, 3, 4))

        with tempfile.TemporaryDirectory() as tmpdir:
            (Path(tmpdir) / "episode_s01e01_2_deu.vtt").write_text(
                "WEBVTT\n\n",
                encoding="utf-8",
            )

            with self.assertRaisesRegex(ValueError, "expected pattern"):
                descriptor.importSubtitles(tmpdir, "episode", strict=True)

    def test_strict_basename_import_rejects_duplicate_track_indices(self):
        descriptor = self.make_descriptor(indices=(2, 3, 4))

        with tempfile.TemporaryDirectory() as tmpdir:
            for filename in ("episode_2_deu.vtt", "episode_2_eng.vtt"):
                (Path(tmpdir) / filename).write_text("WEBVTT\n\n", encoding="utf-8")

            with self.assertRaisesRegex(ValueError, "Multiple external subtitle files"):
                descriptor.importSubtitles(tmpdir, "episode", strict=True)


if __name__ == "__main__":
    unittest.main()