luke_scribe/tests/unit/test_engine_owner.py

"""EngineOwner 단위 테스트 — OOM 강등 체인, attempted_profiles 영속화, 우선순위."""

from __future__ import annotations

import pytest

from luke_scribe.engine.base import TranscriptionOptions
from luke_scribe.engine.owner import MAX_DOWNGRADES, EngineOwner, InferenceRequest
from luke_scribe.errors import OutOfMemory


class OOMFakeEngine:
    """첫 N번 호출에서 OOM, 이후 성공하는 mock."""

    def __init__(self, oom_calls: int = 0) -> None:
        self.oom_calls = oom_calls
        self.calls: list[dict] = []
        self.closed = False

    def transcribe(self, audio_path, options, should_cancel=None, download_progress=None):
        self.calls.append({"device": options.device, "compute_type": options.compute_type})
        if len(self.calls) <= self.oom_calls:
            raise OutOfMemory("CUDA OOM (mock)")
        return type("O", (), {"segments": iter([]), "info": {}})()

    def unload_all(self):
        self.closed = True


def _opts(**kw) -> TranscriptionOptions:
    kw.setdefault("model", "large-v3-turbo")
    kw.setdefault("device", "cuda:0")
    kw.setdefault("compute_type", "float16")
    return TranscriptionOptions(**kw)


class _SegFakeEngine:
    """emit_hypothesis 테스트용 — WAV 경로 기록 + 세그먼트 반환."""

    def __init__(self) -> None:
        self.path = None

    def transcribe(self, audio_path, options, should_cancel=None, download_progress=None):
        self.path = audio_path
        segs = [
            {
                "start": 0.0,
                "end": 1.2,
                "text": "안녕하세요",
                "avg_logprob": -0.2,
                "no_speech_prob": 0.01,
            }
        ]
        return type("O", (), {"segments": iter(segs), "info": {"language": "ko"}})()

    def unload_all(self):
        pass


def _bare_owner(engine) -> EngineOwner:
    import threading

    from luke_scribe.config import Settings

    owner = EngineOwner.__new__(EngineOwner)
    owner.settings = Settings(_env_file=None)
    owner._engine = engine
    owner._lock = threading.Lock()
    owner._realtime_priority = threading.Lock()
    owner._stats = {"jobs": 0, "downgrades": 0, "oom": 0}
    return owner


def test_pcm16_to_wav_header():
    import struct

    from luke_scribe.engine.owner import _pcm16_to_wav

    wav = _pcm16_to_wav(b"\x00\x00" * 8000)  # 1초 (16kHz mono)
    assert wav[:4] == b"RIFF"
    assert wav[8:12] == b"WAVE"
    assert wav[12:16] == b"fmt "
    sr = struct.unpack("<I", wav[24:28])[0]
    channels = struct.unpack("<H", wav[22:24])[0]
    bits = struct.unpack("<H", wav[34:36])[0]
    assert sr == 16000 and channels == 1 and bits == 16
    assert struct.unpack("<I", wav[40:44])[0] == 16000  # data 크기


def test_emit_hypothesis_decodes_chunk_and_cleans_temp():
    import os

    engine = _SegFakeEngine()
    owner = _bare_owner(engine)
    chunk = b"\x00\x00" * 16000  # 1초 PCM16
    out = owner.emit_hypothesis(chunk)
    assert out["audio_sec"] == 1.0
    assert len(out["segments"]) == 1
    assert out["segments"][0].model_dump()["text"] == "안녕하세요"
    # 임시 WAV는 실시간 레인 decode 후 삭제됨
    assert engine.path is not None
    assert not os.path.exists(engine.path)


def test_no_downgrade_on_success():
    owner = EngineOwner.__new__(EngineOwner)
    owner.settings = None
    owner._engine = OOMFakeEngine(oom_calls=0)
    owner._lock = __import__("threading").Lock()
    owner._realtime_priority = __import__("threading").Lock()
    owner._stats = {"jobs": 0, "downgrades": 0, "oom": 0}
    req = InferenceRequest(audio_path="/tmp/x.wav", options=_opts())
    out = owner.transcribe(req)
    assert out["compute_type"] == "float16"
    assert len(out["attempted_profiles"]) == 1


def test_downgrade_chain_on_oom():
    engine = OOMFakeEngine(oom_calls=1)  # 첫 프로파일에서 OOM
    owner = EngineOwner.__new__(EngineOwner)
    owner.settings = None
    owner._engine = engine
    owner._lock = __import__("threading").Lock()
    owner._realtime_priority = __import__("threading").Lock()
    owner._stats = {"jobs": 0, "downgrades": 0, "oom": 0}
    req = InferenceRequest(audio_path="/tmp/x.wav", options=_opts())
    out = owner.transcribe(req)
    # float16 → int8_float16 (강등 1회)
    assert out["compute_type"] == "int8_float16"
    assert len(out["attempted_profiles"]) == 2


def test_downgrade_capped_at_max():
    engine = OOMFakeEngine(oom_calls=99)  # 항상 OOM
    owner = EngineOwner.__new__(EngineOwner)
    owner.settings = None
    owner._engine = engine
    owner._lock = __import__("threading").Lock()
    owner._realtime_priority = __import__("threading").Lock()
    owner._stats = {"jobs": 0, "downgrades": 0, "oom": 0}
    req = InferenceRequest(audio_path="/tmp/x.wav", options=_opts())
    with pytest.raises(OutOfMemory):
        owner.transcribe(req)
    # 강등 체인은 cpu까지 포함하지만 시도는 MAX_DOWNGRADES+1 이내
    assert len(engine.calls) <= MAX_DOWNGRADES + 1


def test_attempted_profiles_persisted_for_retry():
    """재큐(retry) 시 attempted_profiles가 유지되어 총 시도가 캡된다."""
    engine = OOMFakeEngine(oom_calls=0)
    owner = EngineOwner.__new__(EngineOwner)
    owner.settings = None
    owner._engine = engine
    owner._lock = __import__("threading").Lock()
    owner._realtime_priority = __import__("threading").Lock()
    owner._stats = {"jobs": 0, "downgrades": 0, "oom": 0}
    # 이미 2회 시도한 이력이 Job에 영속화된 상황 → 재시도는 즉시 실패해야 함
    req = InferenceRequest(
        audio_path="/tmp/x.wav", options=_opts(), attempted_profiles=["float16", "int8_float16"]
    )
    with pytest.raises(OutOfMemory):
        owner.transcribe(req)
    assert len(engine.calls) == 0  # 새 시도 없음