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

from __future__ import annotations

import json
import subprocess
from pathlib import Path

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

import app.ops.redact as redact_module

from .conftest import AUTH, decode

BOX = {"x": 0.25, "y": 0.25, "w": 0.5, "h": 0.5}
# Pixel region of BOX in the 320x240 fixture.
X0, Y0, X1, Y1 = 80, 60, 240, 180


def _body(media_root: Path, out: Path, mode: str, **extra: object) -> dict[str, object]:
    return {
        "path": f"{media_root}/testsrc.mp4",
        "outPath": str(out),
        "redactions": [
            {
                "keyframes": [{"tMs": 0, "box": BOX}, {"tMs": 2000, "box": BOX}],
                "mode": mode,
                "paddingPct": 0,
                "featherPct": 0,
                **extra,
            }
        ],
    }


def _outside_mask(shape: tuple[int, ...], margin: int) -> np.ndarray:
    mask = np.ones(shape[:2], bool)
    mask[Y0 - margin : Y1 + margin, X0 - margin : X1 + margin] = False
    return mask


def _streams(path: Path) -> list[str]:
    out = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "stream=codec_type,codec_name", "-of", "json", str(path)],
        check=True,
        capture_output=True,
        text=True,
    ).stdout
    return [f"{s['codec_type']}:{s['codec_name']}" for s in json.loads(out)["streams"]]


def test_solid_redaction_changes_only_the_box(client: TestClient, media_root: Path, tmp_path: Path) -> None:
    out = media_root / "redact" / "solid.mkv"
    response = client.post("/v1/redact", headers=AUTH, json=_body(media_root, out, "solid", color="#000000"))
    assert response.status_code == 200, response.text
    assert response.json() == {"outPath": str(out), "framesProcessed": 50}
    source, result = decode(media_root / "testsrc.mp4"), decode(out)
    assert result.shape == source.shape
    inside = result[:, Y0 + 2 : Y1 - 2, X0 + 2 : X1 - 2]
    assert inside.max() <= 8
    assert np.abs(source[:, Y0:Y1, X0:X1].astype(int)).mean() > 50
    outside = _outside_mask(source.shape[1:], 2)
    diff = np.abs(result.astype(int) - source.astype(int))[:, outside]
    assert diff.mean() < 2.0
    assert _streams(out) == ["video:ffv1", "audio:aac"]
    assert not [p for p in out.parent.iterdir() if ".partial-" in p.name]


@pytest.mark.parametrize("mode", ["blur", "pixelate"])
def test_blur_and_pixelate_destroy_detail_inside(client: TestClient, media_root: Path, mode: str) -> None:
    out = media_root / "redact" / f"{mode}.mp4"
    body = _body(media_root, out, mode) | {"path": f"{media_root}/noise.mp4"}
    response = client.post("/v1/redact", headers=AUTH, json=body)
    assert response.status_code == 200, response.text
    source, result = decode(media_root / "noise.mp4"), decode(out)
    region = (slice(None), slice(Y0 + 4, Y1 - 4), slice(X0 + 4, X1 - 4))

    # Local detail (gradient energy) collapses inside the box.
    def detail(frames: np.ndarray) -> float:
        gray = frames.astype(np.float32).mean(axis=3)
        return float(np.abs(np.diff(gray, axis=2)).mean())

    assert detail(result[region]) < detail(source[region]) * 0.2
    assert (
        result[region].astype(np.float32).var(axis=(1, 2)).mean()
        < source[region].astype(np.float32).var(axis=(1, 2)).mean() * 0.2
    )
    outside = _outside_mask(source.shape[1:], 4)
    assert np.abs(result.astype(int) - source.astype(int))[:, outside].mean() < 3.0
    assert _streams(out) == ["video:h264", "audio:aac"]


@pytest.mark.parametrize("mode", ["sticker", "cover", "partial_eyes", "partial_upper", "partial_lower"])
def test_other_modes_produce_output(client: TestClient, media_root: Path, mode: str) -> None:
    out = media_root / "redact" / f"{mode}.mkv"
    body = _body(media_root, out, mode)
    body["redactions"][0]["featherPct"] = 10  # type: ignore[index]
    body["redactions"][0]["paddingPct"] = 20  # type: ignore[index]
    if mode == "partial_eyes":
        body["redactions"][0]["landmarks"] = [  # type: ignore[index]
            {"tMs": 0, "points": [{"x": 0.4, "y": 0.4}, {"x": 0.6, "y": 0.4}, {"x": 0.5, "y": 0.5}]}
        ]
    response = client.post("/v1/redact", headers=AUTH, json=body)
    assert response.status_code == 200, response.text
    source, result = decode(media_root / "testsrc.mp4"), decode(out)
    changed = np.abs(result.astype(int) - source.astype(int))[:, Y0:Y1, X0:X1].mean()
    assert changed > 10
    corner = np.abs(result.astype(int) - source.astype(int))[:, :20, :20].mean()
    assert corner < 2.0


def test_mask_dir_redaction(client: TestClient, media_root: Path) -> None:
    masks = media_root / "redact-masks"
    masks.mkdir(exist_ok=True)
    for t in range(0, 2000, 200):
        mask = np.zeros((240, 320), np.uint8)
        mask[100:140, 140:180] = 255
        cv2.imwrite(str(masks / f"{t}.png"), mask)
    out = media_root / "redact" / "mask.mkv"
    body = {
        "path": f"{media_root}/testsrc.mp4",
        "outPath": str(out),
        "redactions": [
            {
                "keyframes": [],
                "maskDir": str(masks),
                "mode": "cover",
                "color": "#123456",
                "paddingPct": 0,
                "featherPct": 0,
            }
        ],
    }
    response = client.post("/v1/redact", headers=AUTH, json=body)
    assert response.status_code == 200, response.text
    result = decode(out)
    patch = result[:, 104:136, 144:176].reshape(-1, 3).mean(axis=0)
    assert np.allclose(patch, [0x56, 0x34, 0x12], atol=6)


def test_failure_mid_video_leaves_no_output(
    client: TestClient, media_root: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    calls = {"n": 0}
    original = redact_module.apply_redactions

    def flaky(frame, t_ms, plans, frame_ms):  # type: ignore[no-untyped-def]
        calls["n"] += 1
        if calls["n"] == 10:
            raise RuntimeError("simulated frame failure")
        return original(frame, t_ms, plans, frame_ms)

    monkeypatch.setattr(redact_module, "apply_redactions", flaky)
    out = media_root / "redact-fail" / "out.mkv"
    response = client.post("/v1/redact", headers=AUTH, json=_body(media_root, out, "blur"))
    assert response.status_code == 500
    assert response.json()["error"]["code"] == "internal"
    assert not out.exists()
    assert list(out.parent.iterdir()) == []


def test_unreadable_input_fails_closed(client: TestClient, media_root: Path) -> None:
    broken = media_root / "broken.mp4"
    data = (media_root / "testsrc.mp4").read_bytes()
    broken.write_bytes(data[: len(data) // 3])
    out = media_root / "redact-broken" / "out.mkv"
    response = client.post(
        "/v1/redact", headers=AUTH, json=_body(media_root, out, "blur") | {"path": str(broken)}
    )
    assert response.status_code in (400, 500)
    assert not out.exists()
    if out.parent.exists():
        assert not list(out.parent.iterdir())