path: root/tests/test_audio.py

"""Tests for audio helpers: speed parameters, chunk cleanup, encoding,
command construction, duration verification, and audio concatenation."""

import tempfile
import unittest
import wave
from pathlib import Path
from unittest.mock import patch

from converter import audio
from converter import config
from converter.audio import (
    TrackMeta,
    _collect_chunk_files,
    _cover_args,
    _encode_args,
    _tag_args,
    build_concat_command,
    build_ffmetadata,
    build_m4b_chapters_command,
    cleanup_chunks,
    concat_audio_files,
    speed_export_params,
    verify_output_duration,
)


class SpeedExportParamsTests(unittest.TestCase):
    def test_normal_speed_no_filter(self):
        self.assertEqual(speed_export_params(1.0), [])

    def test_simple_speedup(self):
        self.assertEqual(speed_export_params(1.5), ["-filter:a", "atempo=1.5"])

    def test_simple_slowdown(self):
        self.assertEqual(speed_export_params(0.75), ["-filter:a", "atempo=0.75"])

    def test_chained_speedup_beyond_2x(self):
        self.assertEqual(speed_export_params(3.0), ["-filter:a", "atempo=2.0,atempo=1.5"])

    def test_chained_slowdown_below_half(self):
        self.assertEqual(speed_export_params(0.25), ["-filter:a", "atempo=0.5,atempo=0.5"])

    def test_zero_speed_rejected(self):
        with self.assertRaises(ValueError):
            speed_export_params(0)

    def test_negative_speed_rejected(self):
        with self.assertRaises(ValueError):
            speed_export_params(-1.5)


class CleanupChunksTests(unittest.TestCase):
    def test_removes_only_chunk_files(self):
        with tempfile.TemporaryDirectory() as tmp:
            chunks_dir = Path(tmp)
            (chunks_dir / "chunk_0001.wav").write_bytes(b"stale")
            (chunks_dir / "chunk_0002.wav").write_bytes(b"stale")
            (chunks_dir / "keep.txt").write_bytes(b"keep")

            original = audio.CHUNKS_FOLDER
            audio.CHUNKS_FOLDER = chunks_dir
            try:
                cleanup_chunks()
            finally:
                audio.CHUNKS_FOLDER = original

            self.assertFalse((chunks_dir / "chunk_0001.wav").exists())
            self.assertFalse((chunks_dir / "chunk_0002.wav").exists())
            self.assertTrue((chunks_dir / "keep.txt").exists())

    def test_removes_chapter_files(self):
        with tempfile.TemporaryDirectory() as tmp:
            chunks_dir = Path(tmp)
            (chunks_dir / "chapter_0001.m4b").write_bytes(b"stale")
            (chunks_dir / "chunk_0001.wav").write_bytes(b"stale")

            original = audio.CHUNKS_FOLDER
            audio.CHUNKS_FOLDER = chunks_dir
            try:
                cleanup_chunks()
            finally:
                audio.CHUNKS_FOLDER = original

            self.assertFalse((chunks_dir / "chapter_0001.m4b").exists())
            self.assertFalse((chunks_dir / "chunk_0001.wav").exists())


class EncodeArgsTests(unittest.TestCase):
    def test_mp3_uses_bitrate_only(self):
        self.assertEqual(_encode_args("mp3"), ["-b:a", config.AUDIO_BITRATE])

    def test_m4b_uses_aac(self):
        self.assertEqual(_encode_args("m4b"), ["-c:a", "aac", "-b:a", config.AUDIO_BITRATE])

    def test_wav_is_lossless_pcm(self):
        self.assertEqual(_encode_args("wav"), ["-c:a", "pcm_s16le"])

    def test_ogg_uses_libvorbis(self):
        self.assertEqual(_encode_args("ogg"), ["-c:a", "libvorbis", "-b:a", config.AUDIO_BITRATE])

    def test_flac_is_lossless(self):
        self.assertEqual(_encode_args("flac"), ["-c:a", "flac"])


class M4bContainerArgsTests(unittest.TestCase):
    def setUp(self):
        self._original = audio._brand_supported
        audio._brand_supported = True

    def tearDown(self):
        audio._brand_supported = self._original

    def test_includes_faststart_and_brand(self):
        args = audio._m4b_container_args()
        self.assertIn("+faststart", args)
        self.assertIn("M4B ", args)

    def test_brand_omitted_when_unsupported(self):
        audio._brand_supported = False
        self.assertEqual(audio._m4b_container_args(), ["-movflags", "+faststart"])


class BuildConcatCommandTests(unittest.TestCase):
    def setUp(self):
        self._original = audio._brand_supported
        audio._brand_supported = True

    def tearDown(self):
        audio._brand_supported = self._original

    def test_mp3_has_no_container_flags(self):
        cmd = build_concat_command(Path("list.txt"), Path("out.mp3"), "mp3")
        self.assertEqual(cmd[:6], ["ffmpeg", "-y", "-f", "concat", "-safe", "0"])
        self.assertNotIn("-movflags", cmd)
        self.assertEqual(cmd[-1], "out.mp3")

    def test_m4b_gets_faststart_and_brand(self):
        cmd = build_concat_command(Path("list.txt"), Path("out.m4b"), "m4b")
        self.assertIn("+faststart", cmd)
        self.assertIn("M4B ", cmd)
        self.assertEqual(cmd[-1], "out.m4b")

    def test_speed_copy_writes_two_outputs(self):
        cmd = build_concat_command(Path("list.txt"), Path("out.m4b"), "m4b",
                                   speed=1.5, speed_path=Path("out_1.5.m4b"))
        self.assertIn("out.m4b", cmd)
        self.assertIn("out_1.5.m4b", cmd)
        # faststart must apply to both outputs
        self.assertEqual(cmd.count("+faststart"), 2)
        self.assertTrue(any("atempo=1.5" in arg for arg in cmd))

    def test_wav_intermediate(self):
        cmd = build_concat_command(Path("list.txt"), Path("chapter.wav"), "wav")
        self.assertIn("pcm_s16le", cmd)
        self.assertNotIn("-movflags", cmd)

    def test_speed_without_speed_path_rejected(self):
        with self.assertRaises(ValueError):
            build_concat_command(Path("list.txt"), Path("out.mp3"), "mp3", speed=1.5)


class CollectChunkFilesTests(unittest.TestCase):
    def test_uses_recorded_paths_exactly(self):
        with tempfile.TemporaryDirectory() as tmp:
            present = Path(tmp) / "chunk_0001.wav"
            present.write_bytes(b"audio")
            chunk_results = {
                1: present,
                2: None,  # failed chunk
                3: Path(tmp) / "chunk_0003.wav",  # recorded but deleted
            }
            files, missing = _collect_chunk_files(3, chunk_results)

        self.assertEqual(files, [present])
        self.assertEqual(missing, [2, 3])


class BuildM4bChaptersCommandTests(unittest.TestCase):
    def setUp(self):
        self._original = audio._brand_supported
        audio._brand_supported = True

    def tearDown(self):
        audio._brand_supported = self._original

    def test_base_output_maps_metadata_and_chapters(self):
        cmd = build_m4b_chapters_command(Path("list.txt"), Path("meta.txt"), Path("out.m4b"))
        self.assertIn("-map_metadata", cmd)
        self.assertIn("-map_chapters", cmd)
        self.assertIn("out.m4b", cmd)
        self.assertIn("+faststart", cmd)
        self.assertNotIn("filter_complex", cmd)

    def test_speed_outputs_get_their_own_chapter_metadata(self):
        cmd = build_m4b_chapters_command(
            Path("list.txt"), Path("meta.txt"), Path("out.m4b"),
            speed=2.0, speed_path=Path("out_2.m4b"), speed_metadata_file=Path("meta2.txt"),
        )
        self.assertEqual(cmd.count("+faststart"), 2)
        self.assertEqual(cmd.count("-map_chapters"), 2)
        self.assertTrue(any("atempo=2" in arg for arg in cmd))
        # base output chapters come from metadata input 1, speed copy from 2
        chapter_flags = [i for i, v in enumerate(cmd) if v == "-map_chapters"]
        self.assertEqual(cmd[chapter_flags[0] + 1], "1")
        self.assertEqual(cmd[chapter_flags[1] + 1], "2")
        base_idx, speed_idx = cmd.index("out.m4b"), cmd.index("out_2.m4b")
        self.assertLess(chapter_flags[0], base_idx)
        self.assertGreater(chapter_flags[1], base_idx)
        self.assertLess(chapter_flags[1], speed_idx)


class VerifyOutputDurationTests(unittest.TestCase):
    def _patch_probe(self, ms):
        audio.probe_duration_ms = lambda path: ms

    def setUp(self):
        self._original_probe = audio.probe_duration_ms

    def tearDown(self):
        audio.probe_duration_ms = self._original_probe

    def test_close_duration_passes(self):
        self._patch_probe(100_000)
        self.assertTrue(verify_output_duration(Path("x.m4b"), 101_000))

    def test_unverifiable_duration_passes(self):
        self._patch_probe(0)
        self.assertTrue(verify_output_duration(Path("x.m4b"), 100_000))

    def test_zero_expected_passes(self):
        self._patch_probe(50_000)
        self.assertTrue(verify_output_duration(Path("x.m4b"), 0))

    def test_large_drift_fails(self):
        self._patch_probe(3_600_000)  # bogus "1 hour" for a 1 minute book
        with self.assertLogs(level="ERROR"):
            self.assertFalse(verify_output_duration(Path("x.m4b"), 60_000))


class BuildFFMetadataTests(unittest.TestCase):
    def test_writes_chapters(self):
        chapters = [(0, 1200, "One"), (1200, 2500, "Two")]
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "meta.txt"
            build_ffmetadata(chapters, path)
            content = path.read_text(encoding="utf-8")

        self.assertTrue(content.startswith(";FFMETADATA1\n"))
        self.assertIn("[CHAPTER]", content)
        self.assertIn("TIMEBASE=1/1000", content)
        self.assertIn("START=0", content)
        self.assertIn("END=1200", content)
        self.assertIn("title=One", content)
        self.assertIn("START=1200", content)
        self.assertIn("title=Two", content)

    def test_escapes_special_characters(self):
        # ffmpeg's FFMETADATA format treats = ; # and \ as structural.
        chapters = [(0, 1000, "A = B; C# D\\E")]
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "meta.txt"
            build_ffmetadata(chapters, path)
            content = path.read_text(encoding="utf-8")

        self.assertIn(r"title=A \= B\; C\# D\\E", content)

    def test_collapses_newlines_in_titles(self):
        chapters = [(0, 1000, "Two\nLines")]
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "meta.txt"
            build_ffmetadata(chapters, path)
            content = path.read_text(encoding="utf-8")

        self.assertIn("title=Two Lines\n", content)
        self.assertNotIn("title=Two\n", content)


class TagArgsTests(unittest.TestCase):
    META = TrackMeta(title="Dune", artist="Frank Herbert", album="Dune",
                     track=2, total_tracks=5)

    def test_full_meta_written(self):
        args = _tag_args(self.META, "mp3")
        for pair in ("title=Dune", "artist=Frank Herbert",
                     "album=Dune", "track=2/5"):
            self.assertIn(pair, args)

    def test_mp3_gets_id3v23(self):
        mp3_args = _tag_args(self.META, "mp3")
        self.assertIn("-id3v2_version", mp3_args)
        self.assertEqual(mp3_args[mp3_args.index("-id3v2_version") + 1], "3")
        self.assertNotIn("-id3v2_version", _tag_args(self.META, "flac"))

    def test_empty_fields_omitted(self):
        meta = TrackMeta(title="Only Title")
        args = _tag_args(meta, "flac")
        self.assertNotIn("artist", args)
        self.assertNotIn("album", args)
        self.assertNotIn("track", args)

    def test_track_requires_total(self):
        meta = TrackMeta(title="T", track=3)
        self.assertNotIn("track", _tag_args(meta, "mp3"))


class CoverArgsTests(unittest.TestCase):
    def test_mp3_copies_png_stream(self):
        args = _cover_args("mp3", 1)
        self.assertIn("copy", args)
        self.assertIn("attached_pic", args)
        self.assertIn("1:v", args)

    def test_m4b_reencodes_to_jpeg(self):
        args = _cover_args("m4b", 2)
        self.assertIn("mjpeg", args)
        self.assertIn("attached_pic", args)
        self.assertIn("3", args)  # jpeg quality

    def test_ogg_and_wav_have_no_cover(self):
        self.assertEqual(_cover_args("ogg", 1), [])
        self.assertEqual(_cover_args("wav", 1), [])


class BuildConcatCommandMetaTests(unittest.TestCase):
    META = TrackMeta(title="Chapter 1", artist="Author", album="Book",
                     track=1, total_tracks=3)

    def test_cover_added_as_second_input(self):
        cmd = build_concat_command(Path("list.txt"), Path("out.mp3"), "mp3",
                                   meta=self.META, cover=Path("cover.png"))
        # The cover is the second input, after the concat list
        self.assertIn("cover.png", cmd)
        self.assertLess(cmd.index("list.txt"), cmd.index("cover.png"))
        self.assertIn("-map", cmd)
        self.assertIn("1:v", cmd)
        self.assertIn("attached_pic", cmd)
        self.assertEqual(cmd[-1], "out.mp3")

    def test_audio_explicitly_mapped_when_cover_present(self):
        cmd = build_concat_command(Path("list.txt"), Path("out.flac"), "flac",
                                   cover=Path("cover.png"))
        self.assertIn("0:a", cmd)

    def test_no_cover_keeps_single_input(self):
        cmd = build_concat_command(Path("list.txt"), Path("out.mp3"), "mp3",
                                   meta=self.META)
        self.assertEqual(cmd.count("-i"), 1)
        self.assertNotIn("attached_pic", cmd)

    def test_ogg_never_gets_cover_input(self):
        cmd = build_concat_command(Path("list.txt"), Path("out.ogg"), "ogg",
                                   meta=self.META, cover=Path("cover.png"))
        self.assertEqual(cmd.count("-i"), 1)
        self.assertNotIn("attached_pic", cmd)

    def test_speed_copy_gets_tags_and_cover(self):
        cmd = build_concat_command(Path("list.txt"), Path("out.mp3"), "mp3",
                                   speed=1.5, speed_path=Path("out_1.5.mp3"),
                                   meta=self.META, cover=Path("cover.png"))
        self.assertEqual(cmd.count("attached_pic"), 2)
        self.assertEqual(cmd.count("title=Chapter 1"), 2)
        self.assertEqual(cmd.count("1:v"), 2)

    def test_tags_without_cover_present(self):
        cmd = build_concat_command(Path("list.txt"), Path("out.mp3"), "mp3",
                                   meta=self.META)
        self.assertIn("title=Chapter 1", cmd)
        self.assertIn("artist=Author", cmd)
        self.assertIn("album=Book", cmd)
        self.assertIn("track=1/3", cmd)


class BuildM4bChaptersCommandMetaTests(unittest.TestCase):
    def setUp(self):
        self._original = audio._brand_supported
        audio._brand_supported = True

    def tearDown(self):
        audio._brand_supported = self._original

    def test_cover_indexed_after_metadata_inputs(self):
        cmd = build_m4b_chapters_command(Path("list.txt"), Path("meta.txt"),
                                         Path("out.m4b"), cover=Path("cover.png"))
        # Inputs: 0=audio, 1=ffmetadata, 2=cover
        self.assertIn("-i", cmd)
        self.assertIn("2:v", cmd)
        self.assertIn("attached_pic", cmd)

    def test_speed_variant_cover_is_input_three(self):
        cmd = build_m4b_chapters_command(
            Path("list.txt"), Path("meta.txt"), Path("out.m4b"),
            speed=2.0, speed_path=Path("out_2.m4b"), speed_metadata_file=Path("meta2.txt"),
            meta=TrackMeta(title="Book"), cover=Path("cover.png"),
        )
        self.assertEqual(cmd.count("3:v"), 2)  # both outputs attach the cover
        self.assertNotIn("2:v", cmd)
        self.assertEqual(cmd.count("title=Book"), 2)
        # Chapter metadata inputs keep their 1/2 mapping
        chapter_flags = [i for i, v in enumerate(cmd) if v == "-map_chapters"]
        self.assertEqual(cmd[chapter_flags[0] + 1], "1")
        self.assertEqual(cmd[chapter_flags[1] + 1], "2")

    def test_without_cover_regression(self):
        cmd = build_m4b_chapters_command(Path("list.txt"), Path("meta.txt"),
                                         Path("out.m4b"))
        self.assertNotIn("attached_pic", cmd)
        self.assertNotIn("-metadata", cmd)


class ConcatAudioFilesTests(unittest.TestCase):
    """Concatenation of sub-request audio into one chunk file."""

    @staticmethod
    def _write_wav(path: Path, frames: bytes, framerate: int = 24000) -> Path:
        with wave.open(str(path), "wb") as wav_file:
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)
            wav_file.setframerate(framerate)
            wav_file.writeframes(frames)
        return path

    def test_wav_files_are_merged_in_order(self):
        with tempfile.TemporaryDirectory() as tmp:
            first = self._write_wav(Path(tmp) / "a.wav", b"\x01\x00" * 10)
            second = self._write_wav(Path(tmp) / "b.wav", b"\x02\x00" * 20)
            destination = Path(tmp) / "out.wav"
            concat_audio_files([first, second], destination)
            with wave.open(str(destination), "rb") as wav_file:
                self.assertEqual(wav_file.getframerate(), 24000)
                self.assertEqual(wav_file.getnchannels(), 1)
                self.assertEqual(wav_file.getsampwidth(), 2)
                frames = wav_file.readframes(wav_file.getnframes())
            self.assertEqual(frames, b"\x01\x00" * 10 + b"\x02\x00" * 20)

    def test_single_wav_file_is_copied(self):
        with tempfile.TemporaryDirectory() as tmp:
            source = self._write_wav(Path(tmp) / "a.wav", b"\x03\x00" * 15)
            destination = Path(tmp) / "out.wav"
            concat_audio_files([source], destination)
            with wave.open(str(destination), "rb") as wav_file:
                self.assertEqual(wav_file.readframes(wav_file.getnframes()),
                                 b"\x03\x00" * 15)

    def test_empty_source_list_raises(self):
        with tempfile.TemporaryDirectory() as tmp:
            with self.assertRaises(ValueError):
                concat_audio_files([], Path(tmp) / "out.wav")

    def test_mismatched_wav_parameters_fall_back_to_ffmpeg(self):
        with tempfile.TemporaryDirectory() as tmp:
            first = self._write_wav(Path(tmp) / "a.wav", b"\x01\x00" * 10, framerate=24000)
            second = self._write_wav(Path(tmp) / "b.wav", b"\x02\x00" * 10, framerate=16000)
            destination = Path(tmp) / "out.wav"
            with patch("converter.audio.shutil.which", return_value=None), \
                    self.assertRaises(RuntimeError) as ctx:
                concat_audio_files([first, second], destination)
            self.assertIn("ffmpeg", str(ctx.exception))
            # The wave-module path must not have written a partial output.
            self.assertFalse(destination.exists())

    def test_non_wav_input_falls_back_to_ffmpeg(self):
        with tempfile.TemporaryDirectory() as tmp:
            source = Path(tmp) / "part.mp3"
            source.write_bytes(b"not a wav file")
            destination = Path(tmp) / "out.wav"
            with patch("converter.audio.shutil.which", return_value=None), \
                    self.assertRaises(RuntimeError) as ctx:
                concat_audio_files([source], destination)
            self.assertIn("ffmpeg", str(ctx.exception))


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