evOo-local/infra/media-worker/tests/conftest.py

"""Shared fixtures. All media is synthetic (ffmpeg testsrc/sine, drawn shapes)."""

from __future__ import annotations

import subprocess
from collections.abc import Callable
from pathlib import Path

import cv2
import numpy as np
import pytest
import soundfile
from fastapi.testclient import TestClient

from app.config import Settings
from app.main import create_app
from app.models import Backend, Registry, default_backends

from .fakes import fake_backends

TOKEN = "test-token-" + "x" * 32
AUTH = {"Authorization": f"Bearer {TOKEN}"}


def ffmpeg(*args: str) -> None:
    subprocess.run(["ffmpeg", "-v", "error", "-y", *args], check=True)


def write_video(path: Path, frames: list[np.ndarray], fps: int = 25) -> None:
    height, width = frames[0].shape[:2]
    proc = subprocess.Popen(
        [
            "ffmpeg",
            "-v",
            "error",
            "-y",
            "-f",
            "rawvideo",
            "-pix_fmt",
            "bgr24",
            "-s",
            f"{width}x{height}",
            "-r",
            str(fps),
            "-i",
            "pipe:0",
            "-c:v",
            "libx264",
            "-crf",
            "10",
            "-pix_fmt",
            "yuv420p",
            str(path),
        ],
        stdin=subprocess.PIPE,
    )
    assert proc.stdin is not None
    for frame in frames:
        proc.stdin.write(frame.tobytes())
    proc.stdin.close()
    assert proc.wait() == 0


def decode(path: Path) -> np.ndarray:
    probe = (
        subprocess.run(
            [
                "ffprobe",
                "-v",
                "error",
                "-select_streams",
                "v:0",
                "-show_entries",
                "stream=width,height",
                "-of",
                "csv=p=0",
                str(path),
            ],
            check=True,
            capture_output=True,
            text=True,
        )
        .stdout.strip()
        .split(",")
    )
    width, height = int(probe[0]), int(probe[1])
    raw = subprocess.run(
        [
            "ffmpeg",
            "-v",
            "error",
            "-i",
            str(path),
            "-map",
            "0:v:0",
            "-fps_mode",
            "passthrough",
            "-f",
            "rawvideo",
            "-pix_fmt",
            "bgr24",
            "pipe:1",
        ],
        check=True,
        capture_output=True,
    ).stdout
    return np.frombuffer(raw, np.uint8).reshape(-1, height, width, 3)


@pytest.fixture(scope="session")
def media_root(tmp_path_factory: pytest.TempPathFactory) -> Path:
    root = tmp_path_factory.mktemp("media")
    ffmpeg(
        "-f",
        "lavfi",
        "-i",
        "testsrc=size=320x240:rate=25:duration=2",
        "-f",
        "lavfi",
        "-i",
        "sine=frequency=440:duration=2",
        "-c:v",
        "libx264",
        "-crf",
        "10",
        "-pix_fmt",
        "yuv420p",
        "-c:a",
        "aac",
        "-shortest",
        str(root / "testsrc.mp4"),
    )
    ffmpeg(
        "-f",
        "lavfi",
        "-i",
        "testsrc=size=320x240:rate=25:duration=2",
        "-f",
        "lavfi",
        "-i",
        "mandelbrot=size=320x240:rate=25",
        "-filter_complex",
        "[1:v]trim=duration=2,setpts=PTS-STARTPTS[b];[0:v][b]concat=n=2:v=1:a=0",
        "-c:v",
        "libx264",
        "-pix_fmt",
        "yuv420p",
        str(root / "cut.mp4"),
    )
    # Seeded 4-px block noise with a tone: lots of local detail for blur/pixelate checks.
    rng = np.random.default_rng(7)
    noise = [
        np.kron(rng.integers(0, 256, (60, 80, 3), dtype=np.uint8), np.ones((4, 4, 1), np.uint8))
        for _ in range(25)
    ]
    write_video(root / "noise-silent.mp4", noise)
    ffmpeg(
        "-i",
        str(root / "noise-silent.mp4"),
        "-f",
        "lavfi",
        "-i",
        "sine=frequency=330:duration=1",
        "-c:v",
        "copy",
        "-c:a",
        "aac",
        "-shortest",
        str(root / "noise.mp4"),
    )
    # A white disc moving left to right on black: the synthetic "face"/"mark".
    frames = []
    for i in range(40):
        frame = np.zeros((240, 320, 3), np.uint8)
        cv2.circle(frame, (60 + i * 5, 120), 25, (255, 255, 255), -1)
        frames.append(frame)
    write_video(root / "disc.mp4", frames)
    # A bright square in the right half for saliency.
    frame = np.full((240, 320, 3), 40, np.uint8)
    cv2.rectangle(frame, (240, 100), (280, 140), (255, 255, 255), -1)
    write_video(root / "spot.mp4", [frame] * 10)
    # 120 BPM click track (8 s) followed by 1.5 s of silence.
    sr = 22050
    clicks = np.zeros(int(sr * 9.5), np.float32)
    for beat in range(16):
        start = int(beat * 0.5 * sr)
        t = np.arange(int(0.03 * sr)) / sr
        clicks[start : start + t.size] += 0.8 * np.sin(2 * np.pi * 1000 * t) * np.exp(-t * 80)
    soundfile.write(str(root / "clicks.wav"), clicks, sr)
    return root


def make_settings(media_root: Path, tmp_path: Path, **overrides: object) -> Settings:
    values: dict[str, object] = {
        "token": TOKEN,
        "media_root": str(media_root),
        "models_dir": str(tmp_path / "models"),
        "allow_model_download": False,
    }
    values.update(overrides)
    return Settings(**values)  # type: ignore[arg-type]


ClientFactory = Callable[..., TestClient]


@pytest.fixture
def make_client(media_root: Path, tmp_path: Path) -> ClientFactory:
    def make(
        backends: dict[str, Backend] | None = None, fakes: bool = True, **overrides: object
    ) -> TestClient:
        settings = make_settings(media_root, tmp_path, **overrides)
        merged = default_backends()
        if fakes:
            merged.update(fake_backends())
        merged.update(backends or {})
        return TestClient(create_app(settings, Registry(settings, backends=merged)))

    return make


@pytest.fixture
def client(make_client: ClientFactory) -> TestClient:
    return make_client()