tg-favorite-stickers/tests/test_media.py

import asyncio
import json
import shutil
import subprocess
from pathlib import Path

import pytest

from favorite_stickers.media import ConvertedMedia, MediaConversionError, MediaConverter

pytestmark = pytest.mark.skipif(
    not shutil.which("ffmpeg") or not shutil.which("ffprobe"),
    reason="ffmpeg and ffprobe are required",
)


def ffprobe(path: Path) -> dict:
    completed = subprocess.run(
        [
            "ffprobe",
            "-v",
            "error",
            "-show_streams",
            "-show_format",
            "-of",
            "json",
            str(path),
        ],
        check=True,
        capture_output=True,
        text=True,
    )
    return json.loads(completed.stdout)


@pytest.mark.asyncio
async def test_static_image_becomes_valid_video_sticker(tmp_path) -> None:
    source = tmp_path / "source.png"
    await asyncio.to_thread(
        subprocess.run,
        [
            "ffmpeg",
            "-v",
            "error",
            "-f",
            "lavfi",
            "-i",
            "color=c=red:s=320x160",
            "-frames:v",
            "1",
            str(source),
        ],
        check=True,
    )

    converted = await MediaConverter().convert_path(source, "static")
    info = ffprobe(converted.path)
    video = next(
        stream for stream in info["streams"] if stream["codec_type"] == "video"
    )

    assert video["codec_name"] == "vp9"
    assert max(video["width"], video["height"]) == 512
    assert not any(stream["codec_type"] == "audio" for stream in info["streams"])
    assert converted.path.stat().st_size <= 256 * 1024


@pytest.mark.asyncio
async def test_static_image_document_is_auto_detected(tmp_path) -> None:
    source = tmp_path / "source.png"
    await asyncio.to_thread(
        subprocess.run,
        [
            "ffmpeg",
            "-v",
            "error",
            "-f",
            "lavfi",
            "-i",
            "color=c=red:s=320x160",
            "-frames:v",
            "1",
            str(source),
        ],
        check=True,
    )

    converted = await MediaConverter().convert_path(source, "auto")
    info = ffprobe(converted.path)

    assert float(info["format"]["duration"]) >= 0.9


@pytest.mark.asyncio
async def test_video_longer_than_three_seconds_is_rejected(tmp_path) -> None:
    source = tmp_path / "long.mp4"
    await asyncio.to_thread(
        subprocess.run,
        [
            "ffmpeg",
            "-v",
            "error",
            "-f",
            "lavfi",
            "-i",
            "color=c=blue:s=64x64:d=4",
            "-c:v",
            "libx264",
            "-pix_fmt",
            "yuv420p",
            str(source),
        ],
        check=True,
    )

    with pytest.raises(MediaConversionError, match="3"):
        await MediaConverter().convert_path(source, "video")


@pytest.mark.asyncio
async def test_animated_gif_document_is_detected_and_converted(tmp_path) -> None:
    source = tmp_path / "animated.gif"
    await asyncio.to_thread(
        subprocess.run,
        [
            "ffmpeg",
            "-v",
            "error",
            "-f",
            "lavfi",
            "-i",
            "testsrc2=s=128x128:r=10:d=2",
            str(source),
        ],
        check=True,
    )

    converted = await MediaConverter().convert_path(source, "auto")
    info = ffprobe(converted.path)

    assert float(info["format"]["duration"]) <= 3
    assert converted.path.stat().st_size <= 256 * 1024


@pytest.mark.asyncio
async def test_animated_gif_document_longer_than_three_seconds_is_rejected(
    tmp_path,
) -> None:
    source = tmp_path / "long.gif"
    await asyncio.to_thread(
        subprocess.run,
        [
            "ffmpeg",
            "-v",
            "error",
            "-f",
            "lavfi",
            "-i",
            "testsrc2=s=64x64:r=5:d=4",
            str(source),
        ],
        check=True,
    )

    with pytest.raises(MediaConversionError, match="3"):
        await MediaConverter().convert_path(source, "auto")


@pytest.mark.asyncio
async def test_animated_apng_document_is_detected_and_converted(tmp_path) -> None:
    source = tmp_path / "animated.apng"
    await asyncio.to_thread(
        subprocess.run,
        [
            "ffmpeg",
            "-v",
            "error",
            "-f",
            "lavfi",
            "-i",
            "testsrc2=s=128x128:r=10:d=2",
            "-plays",
            "0",
            str(source),
        ],
        check=True,
    )

    converted = await MediaConverter().convert_path(source, "auto")
    info = ffprobe(converted.path)

    assert 1.9 <= float(info["format"]["duration"]) <= 3


@pytest.mark.asyncio
async def test_animated_apng_longer_than_three_seconds_is_rejected(tmp_path) -> None:
    source = tmp_path / "long.apng"
    await asyncio.to_thread(
        subprocess.run,
        [
            "ffmpeg",
            "-v",
            "error",
            "-f",
            "lavfi",
            "-i",
            "testsrc2=s=64x64:r=5:d=4",
            "-plays",
            "0",
            str(source),
        ],
        check=True,
    )

    with pytest.raises(MediaConversionError, match="3"):
        await MediaConverter().convert_path(source, "auto")


@pytest.mark.asyncio
async def test_single_frame_apng_document_is_converted(tmp_path) -> None:
    source = tmp_path / "single.apng"
    await asyncio.to_thread(
        subprocess.run,
        [
            "ffmpeg",
            "-v",
            "error",
            "-f",
            "lavfi",
            "-i",
            "color=c=green:s=64x64",
            "-frames:v",
            "1",
            "-plays",
            "0",
            str(source),
        ],
        check=True,
    )

    converted = await MediaConverter().convert_path(source, "auto")

    assert converted.path.stat().st_size <= 256 * 1024


class SlowConverter(MediaConverter):
    async def _convert_path_unlocked(self, source: Path, kind: str):
        await asyncio.sleep(0.2)
        raise AssertionError("timeout was not enforced")


@pytest.mark.asyncio
async def test_conversion_has_total_timeout(tmp_path) -> None:
    converter = SlowConverter(process_timeout=0.01)

    with pytest.raises(MediaConversionError, match="время"):
        await converter.convert_path(tmp_path / "source.bin", "auto")


class CountingConverter(MediaConverter):
    def __init__(self) -> None:
        super().__init__(max_parallel=2)
        self.active = 0
        self.maximum_active = 0

    async def _convert_path_unlocked(self, source: Path, kind: str):
        self.active += 1
        self.maximum_active = max(self.maximum_active, self.active)
        await asyncio.sleep(0.02)
        self.active -= 1
        return ConvertedMedia(source, "hash")


@pytest.mark.asyncio
async def test_conversion_parallelism_is_bounded(tmp_path) -> None:
    converter = CountingConverter()

    await asyncio.gather(
        *(
            converter.convert_path(tmp_path / f"{index}.bin", "auto")
            for index in range(5)
        )
    )

    assert converter.maximum_active == 2


def test_output_validation_rejects_non_vp9_video() -> None:
    probe = {
        "streams": [
            {
                "codec_type": "video",
                "codec_name": "h264",
                "width": 512,
                "height": 256,
                "avg_frame_rate": "30/1",
            }
        ],
        "format": {"duration": "1.0"},
    }

    with pytest.raises(MediaConversionError, match="VP9"):
        MediaConverter._validate_output(probe)


def test_output_validation_rejects_audio_stream() -> None:
    probe = {
        "streams": [
            {
                "codec_type": "video",
                "codec_name": "vp9",
                "width": 512,
                "height": 256,
                "avg_frame_rate": "30/1",
            },
            {"codec_type": "audio", "codec_name": "opus"},
        ],
        "format": {"duration": "1.0"},
    }

    with pytest.raises(MediaConversionError, match="аудио"):
        MediaConverter._validate_output(probe)


def test_input_validation_rejects_decompression_bomb_dimensions() -> None:
    probe = {
        "streams": [
            {
                "codec_type": "video",
                "codec_name": "png",
                "width": 100_000,
                "height": 100_000,
            }
        ]
    }

    with pytest.raises(MediaConversionError, match="разрешение"):
        MediaConverter._validate_input(probe)


def test_ffmpeg_commands_allow_only_stdin_protocol() -> None:
    converter = MediaConverter()
    source = Path("untrusted.m3u8")
    command = converter._command(source, Path("output.webm"), "video", 32, input_fd=7)

    assert command[command.index("-protocol_whitelist") + 1] == "fd"
    assert command[command.index("-fd") + 1] == "7"
    assert "fd:" in command
    assert str(source) not in command

    probe_command = converter._probe_command(count_frames=False, input_fd=8)
    assert probe_command[probe_command.index("-protocol_whitelist") + 1] == "fd"
    assert probe_command[probe_command.index("-fd") + 1] == "8"
    assert "fd:" in probe_command


@pytest.mark.asyncio
async def test_ffprobe_stderr_is_not_exposed_to_user(tmp_path) -> None:
    ffprobe = tmp_path / "fake-ffprobe"
    ffprobe.write_text(
        "#!/bin/sh\necho '/secret/internal/path: decoder exploded' >&2\nexit 1\n"
    )
    ffprobe.chmod(0o700)
    source = tmp_path / "source.bin"
    source.write_bytes(b"not media")

    with pytest.raises(MediaConversionError) as caught:
        await MediaConverter(ffprobe=str(ffprobe)).convert_path(source, "auto")

    assert "/secret/internal/path" not in str(caught.value)
    assert "decoder exploded" not in str(caught.value)


@pytest.mark.asyncio
async def test_media_manifest_cannot_read_another_local_file(tmp_path) -> None:
    referenced = tmp_path / "referenced.webm"
    await asyncio.to_thread(
        subprocess.run,
        [
            "ffmpeg",
            "-v",
            "error",
            "-f",
            "lavfi",
            "-i",
            "color=red:s=64x64:d=1",
            str(referenced),
        ],
        check=True,
    )
    manifest = tmp_path / "untrusted.ffconcat"
    manifest.write_text(
        f"ffconcat version 1.0\nfile '{referenced.as_posix()}'\n",
        encoding="utf-8",
    )

    with pytest.raises(MediaConversionError, match="не поддерживается"):
        await MediaConverter().convert_path(manifest, "auto")