STT-Assistant/tests/test_output.py

"""Tests for transcript output sinks."""

import json
import wave

import numpy as np
import pytest

from stt_daemon.output import TranscriptSink
from stt_daemon.transcriber import TranscriptResult


def _result(text="hello world"):
    return TranscriptResult(language="en", language_probability=0.98, text=text)


def test_jsonl_includes_level_and_reason(tmp_path):
    sink = TranscriptSink(stdout=False, jsonl_dir=str(tmp_path), jsonl=True)
    sink.emit(_result(), 2.5, rms=0.041234, reason="pause")

    files = list(tmp_path.glob("*.jsonl"))
    assert len(files) == 1

    record = json.loads(files[0].read_text().splitlines()[0])
    assert record["rms"] == round(0.041234, 5)
    assert record["end_reason"] == "pause"
    assert record["duration_s"] == 2.5
    assert record["text"] == "hello world"
    assert record["lang"] == "en"


def test_jsonl_includes_transcribe_seconds(tmp_path):
    sink = TranscriptSink(stdout=False, jsonl_dir=str(tmp_path), jsonl=True)
    sink.emit(_result(), 2.5, rms=0.041234, reason="pause", transcribe_s=0.81234)

    record = json.loads(next(tmp_path.glob("*.jsonl")).read_text().splitlines()[0])
    assert record["transcribe_s"] == round(0.81234, 3)


def test_transcribe_seconds_omitted_when_unknown(tmp_path):
    sink = TranscriptSink(stdout=False, jsonl_dir=str(tmp_path), jsonl=True)
    sink.emit(_result(), 2.5, rms=0.041234, reason="pause")

    record = json.loads(next(tmp_path.glob("*.jsonl")).read_text().splitlines()[0])
    assert "transcribe_s" not in record


def test_jsonl_disabled_writes_nothing(tmp_path, capsys):
    """jsonl=False keeps transcripts off disk even when jsonl_dir is set."""
    sink = TranscriptSink(stdout=True, jsonl_dir=str(tmp_path), jsonl=False)
    sink.emit(_result(), 1.0, rms=0.02, reason="pause")

    assert "hello world" in capsys.readouterr().out
    assert list(tmp_path.glob("*.jsonl")) == []


def test_stdout_includes_level_and_reason(capsys):
    sink = TranscriptSink(stdout=True, jsonl_dir=None)
    sink.emit(_result(), 1.0, rms=0.02, reason="end_of_turn")

    out = capsys.readouterr().out
    assert "rms=0.020" in out
    assert "end_of_turn" in out
    assert "hello world" in out


def test_empty_text_is_skipped(tmp_path, capsys):
    sink = TranscriptSink(stdout=True, jsonl_dir=str(tmp_path), jsonl=True)
    sink.emit(_result(text=""), 1.0, rms=0.02, reason="pause")

    assert capsys.readouterr().out == ""
    assert list(tmp_path.glob("*.jsonl")) == []


def test_saves_wav_when_enabled(tmp_path):
    audio = np.zeros(16000, dtype=np.float32)
    audio[:8000] = 0.5
    logs = tmp_path / "logs"

    sink = TranscriptSink(
        stdout=False,
        jsonl_dir=str(logs),
        jsonl=True,
        save_audio=True,
        audio_dir=str(tmp_path / "audio"),
        sample_rate=16000,
    )
    sink.emit(_result(), 1.0, rms=0.3, reason="pause", audio=audio)

    wavs = list((tmp_path / "audio").rglob("*.wav"))
    assert len(wavs) == 1

    with wave.open(str(wavs[0]), "rb") as wf:
        assert wf.getnchannels() == 1
        assert wf.getsampwidth() == 2
        assert wf.getframerate() == 16000
        assert wf.getnframes() == 16000

    record = json.loads(next(logs.glob("*.jsonl")).read_text().splitlines()[0])
    assert record["audio"] == str(wavs[0])


def test_single_timestamp_per_emit(tmp_path, monkeypatch):
    """All artifacts of one emit() share one timestamp, even across midnight."""
    import stt_daemon.output as output
    from datetime import datetime as real_datetime

    calls = []
    t1 = real_datetime(2026, 9, 20, 23, 59, 59)
    t2 = real_datetime(2026, 9, 21, 0, 0, 1)

    class FakeDatetime:
        @staticmethod
        def now():
            calls.append(1)
            return t1 if len(calls) == 1 else t2

    monkeypatch.setattr(output, "datetime", FakeDatetime)

    logs = tmp_path / "logs"
    sink = TranscriptSink(
        stdout=False,
        jsonl_dir=str(logs),
        jsonl=True,
        save_audio=True,
        audio_dir=str(tmp_path / "audio"),
        sample_rate=16000,
    )
    sink.emit(_result(), 1.0, rms=0.1, reason="pause", audio=np.zeros(160, dtype=np.float32))

    assert len(calls) == 1
    jsonl = next(logs.glob("*.jsonl"))
    assert jsonl.name == "2026-09-20.jsonl"
    record = json.loads(jsonl.read_text().splitlines()[0])
    assert record["ts"].startswith("2026-09-20")
    wav = next((tmp_path / "audio").rglob("*.wav"))
    assert wav.parent.name == "2026-09-20"


def test_no_wav_when_disabled(tmp_path):
    sink = TranscriptSink(
        stdout=False,
        jsonl_dir=None,
        save_audio=False,
        audio_dir=str(tmp_path / "audio"),
        sample_rate=16000,
    )
    sink.emit(_result(), 1.0, rms=0.3, reason="pause", audio=np.zeros(1600, dtype=np.float32))

    assert not (tmp_path / "audio").exists()


def test_save_audio_without_samples_is_safe(tmp_path):
    sink = TranscriptSink(
        stdout=False,
        jsonl_dir=None,
        save_audio=True,
        audio_dir=str(tmp_path / "audio"),
        sample_rate=16000,
    )
    sink.emit(_result(), 1.0, rms=0.3, reason="pause", audio=None)

    assert not (tmp_path / "audio").exists()


def test_failed_audio_save_is_non_fatal(tmp_path, monkeypatch, capsys, caplog):
    import stt_daemon.output as output

    def boom(*args, **kwargs):
        raise OSError("disk full")

    monkeypatch.setattr(output, "write_wav", boom)

    logs = tmp_path / "logs"
    sink = TranscriptSink(
        stdout=True,
        jsonl_dir=str(logs),
        jsonl=True,
        save_audio=True,
        audio_dir=str(tmp_path / "audio"),
        sample_rate=16000,
    )
    with caplog.at_level("WARNING"):
        sink.emit(
            _result(), 1.0, rms=0.3, reason="pause", audio=np.zeros(1600, dtype=np.float32)
        )

    assert "hello world" in capsys.readouterr().out
    record = json.loads(next(logs.glob("*.jsonl")).read_text().splitlines()[0])
    assert record["text"] == "hello world"
    assert "audio" not in record
    assert "Failed to save audio" in caplog.text


def test_failed_serialization_leaves_jsonl_untouched(tmp_path):
    """A record that cannot be serialized must not leave a partial line behind."""
    sink = TranscriptSink(stdout=False, jsonl_dir=str(tmp_path), jsonl=True)
    sink.emit(_result(), 2.5, rms=0.04, reason="pause")
    path = next(tmp_path.glob("*.jsonl"))
    before = path.read_text()

    bad = TranscriptResult(language="en", language_probability=np.float32(0.98), text="boom")
    with pytest.raises(TypeError):
        sink.emit(bad, 2.5, rms=0.04, reason="pause")

    assert path.read_text() == before


def test_jsonl_includes_speaker_fields(tmp_path):
    sink = TranscriptSink(stdout=False, jsonl_dir=str(tmp_path), jsonl=True)
    sink.emit(
        _result(),
        3.2,
        rms=0.031,
        reason="pause",
        seg_id=142,
        speaker_id="spk_0001",
        speaker_status="matched",
        speaker="Efren",
        speaker_score=0.712345,
    )

    record = json.loads(next(tmp_path.glob("*.jsonl")).read_text().splitlines()[0])
    assert record["seg_id"] == 142
    assert record["speaker_id"] == "spk_0001"
    assert record["speaker_status"] == "matched"
    assert record["speaker"] == "Efren"
    assert record["speaker_score"] == round(0.712345, 4)


def test_jsonl_speaker_fields_default_to_nulls(tmp_path):
    sink = TranscriptSink(stdout=False, jsonl_dir=str(tmp_path), jsonl=True)
    sink.emit(_result(), 2.5, rms=0.04, reason="pause")

    record = json.loads(next(tmp_path.glob("*.jsonl")).read_text().splitlines()[0])
    assert record["seg_id"] is None
    assert record["speaker_id"] is None
    assert record["speaker_status"] == "none"
    assert record["speaker"] is None
    assert record["speaker_score"] is None


def test_jsonl_speaker_score_rounded_to_4dp(tmp_path):
    sink = TranscriptSink(stdout=False, jsonl_dir=str(tmp_path), jsonl=True)
    sink.emit(_result(), 1.0, rms=0.02, reason="pause", speaker_score=0.123456)

    record = json.loads(next(tmp_path.glob("*.jsonl")).read_text().splitlines()[0])
    assert record["speaker_score"] == round(0.123456, 4)