tests/unit/test_node4_extractor.py

"""
Unit tests for Node 4: Clip Signal Extractor
(src/envs/subenv2/node4_clip_extractor.py)

OpenCV VideoCapture and FaceLandmarker initialization are mocked so no real
video files or model assets are required.
"""

from __future__ import annotations

import json
import types
from pathlib import Path
from unittest.mock import MagicMock, patch

import numpy as np
import pytest

from src.schemas.subenv2 import ClipSignalObservation


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_RNG = np.random.default_rng(42)


def _random_frame(height: int = 72, width: int = 128) -> np.ndarray:
    return _RNG.integers(0, 256, (height, width, 3), dtype=np.uint8)


def mock_capture(frame_count: int, height: int = 72, width: int = 128) -> MagicMock:
    frames = [_random_frame(height, width) for _ in range(frame_count)]
    read_returns = [(True, f) for f in frames] + [(False, None)]
    read_iter = iter(read_returns)

    cap = MagicMock()
    cap.isOpened.return_value = True
    cap.read.side_effect = lambda: next(read_iter)
    cap.release.return_value = None
    return cap


def _make_landmarker_mock() -> MagicMock:
    lm = types.SimpleNamespace(x=0.5, y=0.5, z=0.0)
    landmarks = [lm] * 478
    detect_result = types.SimpleNamespace(face_landmarks=[landmarks])

    landmarker = MagicMock()
    landmarker.detect.return_value = detect_result
    landmarker.close.return_value = None
    return landmarker


def _full_patch(cap_mock: MagicMock, landmarker: MagicMock):
    from contextlib import ExitStack

    stack = ExitStack()
    stack.enter_context(patch("cv2.VideoCapture", return_value=cap_mock))
    stack.enter_context(
        patch(
            "src.envs.subenv2.node4_clip_extractor._create_face_landmarker",
            return_value=landmarker,
        )
    )
    return stack


_EMPTY_CTX: dict = {
    "current_phoneme_coverage": {},
    "current_pose_distribution": {},
    "clips_audited_so_far": 0,
    "similar_clips_accepted": 0,
}


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


def test_raises_on_short_clip(tmp_path):
    dummy = tmp_path / "dummy.mp4"
    dummy.touch()

    cap = mock_capture(10)
    with patch("cv2.VideoCapture", return_value=cap):
        with pytest.raises(ValueError, match="24"):
            from src.envs.subenv2.node4_clip_extractor import extract_clip_signals

            extract_clip_signals(dummy, {})


def test_blur_score_in_range(tmp_path):
    dummy = tmp_path / "dummy.mp4"
    dummy.touch()

    cap30 = mock_capture(30)
    landmarker = _make_landmarker_mock()

    with _full_patch(cap30, landmarker):
        from src.envs.subenv2.node4_clip_extractor import extract_clip_signals

        result = extract_clip_signals(dummy, _EMPTY_CTX)

    assert isinstance(result, ClipSignalObservation)
    assert 0.0 <= result.blur_score <= 1.0


def test_phoneme_coverage_new_empty_dataset(tmp_path):
    dummy = tmp_path / "dummy.mp4"
    dummy.touch()

    aligner_data = {"phonemes": ["AH", "EE", "OW"]}
    aligner_json = tmp_path / "align.json"
    aligner_json.write_text(json.dumps(aligner_data))

    cap30 = mock_capture(30)
    landmarker = _make_landmarker_mock()
    ctx = {**_EMPTY_CTX, "current_phoneme_coverage": {}}

    with _full_patch(cap30, landmarker):
        from src.envs.subenv2.node4_clip_extractor import extract_clip_signals

        result = extract_clip_signals(
            dummy,
            ctx,
            aligner_output=json.loads(aligner_json.read_text()),
        )

    assert result.phoneme_coverage_new == 1.0


def test_phoneme_coverage_new_partial(tmp_path):
    dummy = tmp_path / "dummy.mp4"
    dummy.touch()

    aligner_data = {"phonemes": ["AH", "EE", "OW"]}
    aligner_json = tmp_path / "align.json"
    aligner_json.write_text(json.dumps(aligner_data))

    cap30 = mock_capture(30)
    landmarker = _make_landmarker_mock()
    ctx = {**_EMPTY_CTX, "current_phoneme_coverage": {"AH": 3}}

    with _full_patch(cap30, landmarker):
        from src.envs.subenv2.node4_clip_extractor import extract_clip_signals

        result = extract_clip_signals(
            dummy,
            ctx,
            aligner_output=json.loads(aligner_json.read_text()),
        )

    assert abs(result.phoneme_coverage_new - 2 / 3) < 1e-6


def test_no_forced_align_path(tmp_path):
    dummy = tmp_path / "dummy.mp4"
    dummy.touch()

    cap30 = mock_capture(30)
    landmarker = _make_landmarker_mock()

    with _full_patch(cap30, landmarker):
        from src.envs.subenv2.node4_clip_extractor import extract_clip_signals

        result = extract_clip_signals(dummy, _EMPTY_CTX, aligner_output=None)

    assert result.phoneme_sequence == []
    assert result.lip_sync_confidence == 0.0


def test_raises_when_landmarker_unavailable(tmp_path):
    dummy = tmp_path / "dummy.mp4"
    dummy.touch()

    cap30 = mock_capture(30)
    with patch("cv2.VideoCapture", return_value=cap30), patch(
        "src.envs.subenv2.node4_clip_extractor._create_face_landmarker",
        return_value=None,
    ):
        with pytest.raises(ValueError, match="FaceLandmarker model file"):
            from src.envs.subenv2.node4_clip_extractor import extract_clip_signals

            extract_clip_signals(dummy, _EMPTY_CTX)