tests/test_diarization.py

"""話者分離の呼び出し口。

pyannote は 3 系と 4 系で引数名も戻り値の形も違い、4 系は音声の読み込みに
torchcodec を要求する。モデル本体は動かさず、その差を吸収できているかを確かめる。
"""

import struct
import sys
import wave
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import pipeline as pl  # noqa: E402


class FakeTurn:
    def __init__(self, start, end):
        self.start, self.end = start, end


class FakeAnnotation:
    """pyannote 3 系がそのまま返してくるもの。"""

    def __init__(self, rows):
        self.rows = rows

    def itertracks(self, yield_label=False):
        for start, end, label in self.rows:
            yield FakeTurn(start, end), None, label


class FakeDiarizeOutput:
    """pyannote 4 系が返す入れ物。中に Annotation が入っている。"""

    def __init__(self, annotation):
        self.speaker_diarization = annotation


ROWS = [(0.0, 8.5, "SPEAKER_00"), (8.5, 14.5, "SPEAKER_01")]


@pytest.fixture
def wav(tmp_path) -> Path:
    """16kHz モノラル 16bit の、1秒ぶんの WAV。"""
    path = tmp_path / "sample.wav"
    with wave.open(str(path), "wb") as f:
        f.setnchannels(1)
        f.setsampwidth(2)
        f.setframerate(16000)
        f.writeframes(b"".join(struct.pack("<h", (i % 200) * 100 - 10000) for i in range(16000)))
    return path


def test_WAVを自前で読める(wav):
    """torchcodec が無くても音声を渡せること。"""
    data = pl.load_waveform(wav)
    assert data["sample_rate"] == 16000
    assert tuple(data["waveform"].shape) == (1, 16000)      # (チャンネル, 時間)
    assert -1.0 <= float(data["waveform"].min()) <= float(data["waveform"].max()) <= 1.0


@pytest.mark.parametrize(
    "result",
    [FakeAnnotation(ROWS), FakeDiarizeOutput(FakeAnnotation(ROWS))],
    ids=["pyannote3", "pyannote4"],
)
def test_どちらの戻り値でも同じ区間になる(monkeypatch, wav, result):
    monkeypatch.setattr(pl, "load_diarizer", lambda token, device="": lambda audio, **kw: result)
    assert pl.diarize(wav, "token", None) == ROWS


def test_パスではなく波形を渡している(monkeypatch, wav):
    """パスを渡すと pyannote 4 が torchcodec を要求して落ちるため。"""
    seen = {}

    def fake_pipe(audio, **kwargs):
        seen["audio"] = audio
        seen["kwargs"] = kwargs
        return FakeAnnotation(ROWS)

    monkeypatch.setattr(pl, "load_diarizer", lambda token, device="": fake_pipe)
    pl.diarize(wav, "token", 2)

    assert set(seen["audio"]) == {"waveform", "sample_rate"}
    assert seen["kwargs"] == {"num_speakers": 2}


def test_話者人数を指定しなければ自動判定に任せる(monkeypatch, wav):
    seen = {}
    monkeypatch.setattr(
        pl, "load_diarizer",
        lambda token, device="": lambda audio, **kw: (seen.update(kw), FakeAnnotation(ROWS))[1],
    )
    pl.diarize(wav, "token", None)
    assert seen == {}