scripts/test_english_e2e.py

#!/usr/bin/env python3
"""영어 E2E 파이프라인 테스트.

RAVDESS WAV + 수동 작성 텍스트로 Stage2 전체 파이프라인을 테스트한다.
감정별 2개씩 14개 세그먼트를 통과시켜 audio + text + fusion 동작을 검증.

Usage:
    python scripts/test_english_e2e.py
"""

from __future__ import annotations

import csv
import json
import logging
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("test_english_e2e")

MANIFEST_PATH = PROJECT_ROOT / "data" / "ravdess" / "manifest.csv"

# 감정별 매칭 텍스트 (RAVDESS는 텍스트 없으므로 수동 작성)
EMOTION_TEXTS = {
    "neutral": [
        "The meeting is scheduled for three o'clock.",
        "I'll pick up the groceries on the way home.",
    ],
    "joy": [
        "I'm so happy to hear from you!",
        "That's wonderful news, I'm thrilled!",
    ],
    "sadness": [
        "I miss you so much, it hurts.",
        "I feel really down today, nothing is going right.",
    ],
    "anger": [
        "I can't believe you did that, I'm so angry!",
        "This is completely unacceptable, stop it now!",
    ],
    "surprise": [
        "Oh my god, I can't believe it!",
        "What?! I never expected that!",
    ],
    "fear": [
        "I'm scared, something doesn't feel right.",
        "Please help, I'm terrified right now.",
    ],
    "disgust": [
        "That's absolutely disgusting, I feel sick.",
        "This is revolting, I can't stand it.",
    ],
}


def pick_samples(manifest_path: Path) -> list[dict]:
    """감정별 2개씩 WAV 선택."""
    by_emotion: dict[str, list[dict]] = {}
    with open(manifest_path) as f:
        for row in csv.DictReader(f):
            emotion = row["emotion"]
            if emotion not in by_emotion:
                by_emotion[emotion] = []
            if len(by_emotion[emotion]) < 2:
                by_emotion[emotion].append(row)

    samples = []
    for emotion in EMOTION_TEXTS:
        if emotion in by_emotion:
            samples.extend(by_emotion[emotion])
    return samples


def run_e2e():
    """E2E 테스트 실행."""
    from src.common.schemas import (
        Models,
        ProcessingInfo,
        Segment,
        Stage1Output,
    )
    from src.stage2.process import process

    if not MANIFEST_PATH.exists():
        logger.error("manifest.csv를 찾을 수 없습니다. 먼저 prepare_ravdess.py를 실행하세요.")
        return False

    samples = pick_samples(MANIFEST_PATH)
    logger.info(f"선택된 테스트 샘플: {len(samples)}개")

    # Stage1Output 생성
    segments = []
    ground_truths = []
    for i, sample in enumerate(samples):
        emotion = sample["emotion"]
        texts = EMOTION_TEXTS[emotion]
        text = texts[i % len(texts)]

        segments.append(Segment(
            segment_id=i,
            speaker_id=f"speaker_{int(sample['actor_id']) % 2}",
            start=float(i * 3.0),
            end=float(i * 3.0 + 3.0),
            text=text,
            language="en",
            audio_path=sample["clean_path"],
            confidence=0.95,
        ))
        ground_truths.append(emotion)

    stage1_output = Stage1Output(
        call_id="ravdess_e2e_test",
        duration=float(len(segments) * 3.0),
        speakers=["speaker_0", "speaker_1"],
        audio_path="data/ravdess/test_call.wav",
        segments=segments,
        processing_info=ProcessingInfo(
            processing_time_sec=1.0,
            models=Models(
                diarization="pyannote/speaker-diarization-3.1",
                asr="large-v3-turbo",
                language_id="whisper",
                alignment="whisperx",
            ),
            device="cpu",
        ),
    )

    # Stage2 config (영어 전용)
    config = {
        "audio_emotion": {
            "model": "iic/emotion2vec_plus_base",
        },
        "text_emotion": {
            "korean_model": "searle-j/kote_for_easygoing_people",
            "english_model": "j-hartmann/emotion-english-distilroberta-base",
        },
        "fusion": {
            "audio_weight": 0.6,
            "text_weight": 0.4,
        },
        "output_path": "data/e2e_test_output.json",
    }

    logger.info("Stage 2 process() 실행 중...")
    output = process(stage1_output, config=config)

    # 검증
    logger.info("\n" + "=" * 70)
    logger.info("E2E 테스트 결과")
    logger.info("=" * 70)

    errors = []

    # 1. EmotionResult 개수 확인
    if len(output.emotions) != len(segments):
        errors.append(f"EmotionResult 개수 불일치: {len(output.emotions)} != {len(segments)}")
    else:
        logger.info(f"[PASS] EmotionResult 개수: {len(output.emotions)}")

    # 2. SpeakerSummary 확인
    if len(output.speaker_summaries) > 0:
        logger.info(f"[PASS] SpeakerSummary 생성: {len(output.speaker_summaries)}명")
    else:
        errors.append("SpeakerSummary가 비어있음")

    # 3. 각 EmotionResult 유효성 확인
    valid_emotions = {"neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"}
    for em in output.emotions:
        if em.fused_emotion not in valid_emotions:
            errors.append(f"유효하지 않은 fused_emotion: {em.fused_emotion}")
        if not (0.0 <= em.fused_confidence <= 1.0):
            errors.append(f"fused_confidence 범위 초과: {em.fused_confidence}")

    if not any("유효하지 않은" in e for e in errors):
        logger.info("[PASS] 모든 EmotionResult 유효")

    # 4. JSON 직렬화 테스트
    try:
        json_str = output.model_dump_json(indent=2)
        roundtrip = json.loads(json_str)
        logger.info(f"[PASS] JSON 직렬화/역직렬화 성공 ({len(json_str)} bytes)")
    except Exception as e:
        errors.append(f"JSON 직렬화 실패: {e}")

    # 5. 결과 테이블
    logger.info(f"\n{'Seg':>3s} | {'Ground Truth':>12s} | {'Audio':>10s} | {'Text':>10s} | {'Fused':>10s} | {'Conf':>5s}")
    logger.info("-" * 65)
    audio_correct = 0
    text_correct = 0
    fused_correct = 0
    for i, (em, gt) in enumerate(zip(output.emotions, ground_truths)):
        a_match = "o" if em.audio_emotion == gt else "x"
        t_match = "o" if em.text_emotion == gt else "x"
        f_match = "o" if em.fused_emotion == gt else "x"
        if em.audio_emotion == gt:
            audio_correct += 1
        if em.text_emotion == gt:
            text_correct += 1
        if em.fused_emotion == gt:
            fused_correct += 1
        logger.info(
            f"{i:3d} | {gt:>12s} | {em.audio_emotion:>8s} {a_match} | "
            f"{em.text_emotion:>8s} {t_match} | {em.fused_emotion:>8s} {f_match} | "
            f"{em.fused_confidence:.3f}"
        )

    n = len(ground_truths)
    logger.info(f"\nAccuracy: audio={audio_correct}/{n} text={text_correct}/{n} fused={fused_correct}/{n}")

    # SpeakerSummary 출력
    logger.info("\nSpeaker Summaries:")
    for spk, summary in output.speaker_summaries.items():
        logger.info(f"  {spk}: dominant={summary.dominant_emotion}, "
                     f"conf={summary.avg_confidence:.3f}, "
                     f"dist={summary.emotion_distribution}")

    if errors:
        logger.error(f"\nFAILED — {len(errors)} errors:")
        for e in errors:
            logger.error(f"  - {e}")
        return False

    logger.info(f"\n{'='*70}")
    logger.info("E2E 테스트 PASS — 영어 파이프라인 정상 동작")
    logger.info(f"{'='*70}")
    return True


if __name__ == "__main__":
    success = run_e2e()
    sys.exit(0 if success else 1)