STT-Assistant/tests/test_segmenter.py

"""Tests for the VAD segmenter with pre-roll buffer."""

import numpy as np
import pytest
import webrtcvad

from stt_daemon.segmenter import FrameSegmenter, make_vad


class FakeVad:
    """A VAD for testing: returns scripted speech/non-speech decisions."""

    def __init__(self, decisions: list[bool]):
        """Initialize with a sequence of boolean decisions.

        Args:
            decisions: List of True (speech) / False (silence) per frame.
        """
        self.decisions = decisions
        self.frame_index = 0

    def is_speech(self, frame: bytes, sample_rate: int) -> bool:
        """Return the next scripted decision."""
        if self.frame_index >= len(self.decisions):
            # If we've exhausted decisions, default to False
            return False
        result = self.decisions[self.frame_index]
        self.frame_index += 1
        return result


def make_frame(samples: np.ndarray) -> bytes:
    """Convert samples to int16 PCM bytes."""
    assert samples.dtype == np.int16
    return samples.tobytes()


def make_silent_frame(frame_samples: int) -> bytes:
    """Create a silent frame (zeros)."""
    return make_frame(np.zeros(frame_samples, dtype=np.int16))


def make_speech_frame(frame_samples: int, value: int = 1000) -> bytes:
    """Create a frame with some amplitude."""
    return make_frame(np.full(frame_samples, value, dtype=np.int16))


def test_make_vad_invalid_aggressiveness():
    """Test that make_vad validates aggressiveness."""
    with pytest.raises(ValueError):
        make_vad(4)
    with pytest.raises(ValueError):
        make_vad(-1)


def test_make_vad_valid():
    """Test that make_vad accepts 0-3."""
    for aggr in (0, 1, 2, 3):
        vad = make_vad(aggr)
        assert isinstance(vad, webrtcvad.Vad)


def test_preroll_preserved():
    """Test that pre-roll buffer content is included in the emitted utterance.

    Pre-roll frames are kept in a ring buffer before speech is confirmed.
    When speech is confirmed, those pre-roll frames should be at the start
    of the emitted utterance, ensuring the first phoneme is not cut off.
    """
    sample_rate = 16000
    frame_ms = 30
    frame_samples = round(sample_rate * frame_ms / 1000)

    # Pre-roll frames: 2 silent frames (60ms silence before speech)
    # Then: 3 speech frames to confirm speech (90ms to trigger confirmation)
    # Then: 1 more speech frame + 2 silence frames to finalize

    fake_vad = FakeVad(
        [
            False,  # Frame 0: pre-roll silence
            False,  # Frame 1: pre-roll silence
            True,   # Frame 2: speech (count=1)
            True,   # Frame 3: speech (count=2)
            True,   # Frame 4: speech (count=3, confirmed)
            True,   # Frame 5: speech (during speaking)
            False,  # Frame 6: silence (count=1)
            False,  # Frame 7: silence (count=2, finalize)
        ]
    )

    segmenter = FrameSegmenter(
        sample_rate=sample_rate,
        frame_ms=frame_ms,
        vad=fake_vad,
        pre_roll_ms=150,   # 5 frames (enough for 2 silence + 3 speech to not lose frame 0)
        speech_start_ms=90,  # 3 frames needed
        silence_ms=60,    # 2 frames to finalize
        min_utterance_ms=30,
        max_utterance_s=30,
    )

    pre_roll_frame_0 = make_silent_frame(frame_samples)
    pre_roll_frame_1 = make_silent_frame(frame_samples)
    speech_frame_2 = make_speech_frame(frame_samples)
    speech_frame_3 = make_speech_frame(frame_samples)
    speech_frame_4 = make_speech_frame(frame_samples)
    speech_frame_5 = make_speech_frame(frame_samples)
    silence_frame_6 = make_silent_frame(frame_samples)
    silence_frame_7 = make_silent_frame(frame_samples)

    # Push pre-roll frames (should be buffered, not emitted)
    assert segmenter.push(pre_roll_frame_0) is None
    assert segmenter.push(pre_roll_frame_1) is None

    # Push speech frames until confirmation (frame 4 confirms)
    assert segmenter.push(speech_frame_2) is None
    assert segmenter.push(speech_frame_3) is None
    utt = segmenter.push(speech_frame_4)
    assert utt is None  # _begin() returns None, not the utterance yet

    # Continue with more frames
    assert segmenter.push(speech_frame_5) is None
    assert segmenter.push(silence_frame_6) is None

    # Final silence triggers finalization
    utt = segmenter.push(silence_frame_7)
    assert utt is not None, "Utterance should be emitted after silence"

    # Verify that the utterance contains the pre-roll frames and speech frames
    # Expected: pre_roll_0, pre_roll_1, speech_2-5, silence_6-7
    # (silence frames kept as tail when finalizing with keep_tail=True)
    expected_sample_count = frame_samples * 8
    actual_sample_count = len(utt.audio)

    assert actual_sample_count == expected_sample_count, (
        f"Expected {expected_sample_count} samples (8 frames), "
        f"got {actual_sample_count}"
    )

    # Verify the audio starts with zeros (the pre-roll silence frames 0-1)
    pre_roll_audio = utt.audio[:2 * frame_samples]
    assert np.allclose(pre_roll_audio, 0, atol=0.01), (
        "Pre-roll should be silent at the start (first 2 frames)"
    )

    # Verify speech frames come next (non-zero)
    speech_audio = utt.audio[2 * frame_samples : 6 * frame_samples]
    assert not np.allclose(speech_audio, 0, atol=0.01), (
        "Speech frames should be non-silent"
    )


def test_silence_finalization():
    """Test that silence_frames consecutive frames trigger finalization."""
    sample_rate = 16000
    frame_ms = 30
    frame_samples = round(sample_rate * frame_ms / 1000)

    fake_vad = FakeVad(
        [
            True,   # Start of speech
            True,
            True,
            True,   # 4 speech frames
            False,  # Silence count=1
            False,  # Silence count=2
            False,  # Silence count=3 (finalizes if silence_ms=90)
        ]
    )

    segmenter = FrameSegmenter(
        sample_rate=sample_rate,
        frame_ms=frame_ms,
        vad=fake_vad,
        pre_roll_ms=60,      # 2 frames
        speech_start_ms=30,  # 1 frame to confirm
        silence_ms=60,       # 2 frames to finalize
        min_utterance_ms=30,
        max_utterance_s=30,
    )

    # Push 1 speech frame to trigger confirmation
    utt = segmenter.push(make_speech_frame(frame_samples))
    assert utt is None

    # Push 3 more speech frames
    for _ in range(3):
        assert segmenter.push(make_speech_frame(frame_samples)) is None

    # Push 1 silence frame (count=1)
    assert segmenter.push(make_silent_frame(frame_samples)) is None

    # Push 2nd silence frame (count=2, finalizes)
    utt = segmenter.push(make_silent_frame(frame_samples))
    assert utt is not None
    # At finalization with 2 trailing silence frames, we keep those as tail.
    # Segment progression:
    # - After frame 0 (speech confirm): segment = [frame_0]
    # - After frames 1-3: segment = [0, 1, 2, 3]
    # - After frame 4 (1st silence): segment = [0, 1, 2, 3, 4], silence_run=1
    # - After frame 5 (2nd silence): segment = [0, 1, 2, 3, 4, 5], silence_run=2
    # Finalize: emit 6 frames, keep tail [4, 5] for next pre-roll
    assert len(utt.audio) == frame_samples * 6


def test_min_utterance_discarded():
    """Test that utterances shorter than min_utterance_ms are discarded."""
    sample_rate = 16000
    frame_ms = 30
    frame_samples = round(sample_rate * frame_ms / 1000)

    fake_vad = FakeVad(
        [
            True,   # 1 speech frame (30ms, less than min 300ms)
            False,  # Silence to finalize
            False,
        ]
    )

    segmenter = FrameSegmenter(
        sample_rate=sample_rate,
        frame_ms=frame_ms,
        vad=fake_vad,
        pre_roll_ms=0,
        speech_start_ms=30,  # 1 frame
        silence_ms=60,       # 2 frames
        min_utterance_ms=300,  # Only keep if >= 300ms (10 frames at 30ms)
        max_utterance_s=30,
    )

    # Confirm speech with 1 frame
    segmenter.push(make_speech_frame(frame_samples))

    # Finalize with 2 silence frames
    segmenter.push(make_silent_frame(frame_samples))
    utt = segmenter.push(make_silent_frame(frame_samples))

    assert utt is None, "Short utterance should be discarded"


def test_max_utterance_forces_finalization():
    """Test that max_utterance_s force-finalizes long continuous speech."""
    sample_rate = 16000
    frame_ms = 30
    frame_samples = round(sample_rate * frame_ms / 1000)
    max_utterance_s = 0.3  # 300ms, ~10 frames

    # Create a long sequence of speech with no silence
    num_frames = 15
    fake_vad = FakeVad([True] * num_frames)

    segmenter = FrameSegmenter(
        sample_rate=sample_rate,
        frame_ms=frame_ms,
        vad=fake_vad,
        pre_roll_ms=0,
        speech_start_ms=30,  # 1 frame to confirm
        silence_ms=60,
        min_utterance_ms=30,
        max_utterance_s=max_utterance_s,
    )

    # Confirm speech
    segmenter.push(make_speech_frame(frame_samples))

    # Push more frames; should finalize when max is hit
    utt = None
    for i in range(1, num_frames):
        utt = segmenter.push(make_speech_frame(frame_samples))
        if utt is not None:
            break

    assert utt is not None, "Max duration should trigger finalization"
    # Should have finalized at ~max_utterance_s (a bit flexible due to frame alignment)
    assert len(utt.audio) >= frame_samples * 10


def test_flush_finalizes_inprogress():
    """Test that flush() finalizes an in-progress utterance."""
    sample_rate = 16000
    frame_ms = 30
    frame_samples = round(sample_rate * frame_ms / 1000)

    fake_vad = FakeVad([False, False, True, True])

    segmenter = FrameSegmenter(
        sample_rate=sample_rate,
        frame_ms=frame_ms,
        vad=fake_vad,
        pre_roll_ms=60,       # 2 frames
        speech_start_ms=60,   # 2 frames
        silence_ms=300,
        min_utterance_ms=30,
        max_utterance_s=30,
    )

    # Pre-roll frames (silence)
    segmenter.push(make_silent_frame(frame_samples))
    segmenter.push(make_silent_frame(frame_samples))

    # Confirm speech (need 2 consecutive)
    segmenter.push(make_speech_frame(frame_samples))
    segmenter.push(make_speech_frame(frame_samples))

    # Now call flush without sending silence
    utt = segmenter.flush()
    assert utt is not None
    # Ring at confirmation has the 2 speech frames, so segment = [speech_0, speech_1]
    assert len(utt.audio) == frame_samples * 2  # The 2 speech frames


def test_no_utterance_if_never_speaks():
    """Test that a stream of only silence emits no utterance."""
    sample_rate = 16000
    frame_ms = 30
    frame_samples = round(sample_rate * frame_ms / 1000)

    fake_vad = FakeVad([False, False, False, False])

    segmenter = FrameSegmenter(
        sample_rate=sample_rate,
        frame_ms=frame_ms,
        vad=fake_vad,
        pre_roll_ms=60,
        speech_start_ms=60,
        silence_ms=60,
        min_utterance_ms=30,
        max_utterance_s=30,
    )

    # Push only silence
    for _ in range(4):
        utt = segmenter.push(make_silent_frame(frame_samples))
        assert utt is None

    # Flush should also return None
    utt = segmenter.flush()
    assert utt is None


def test_blip_discarded_with_realistic_pre_roll():
    """Regression: min_utterance_ms must count real speech, not pre-roll padding.

    With the default 400 ms pre-roll the segment is always at least 13 frames
    long, so a length-based check would never discard a blip. Only the actual
    number of speech frames should decide.
    """
    sample_rate = 16000
    frame_ms = 30
    frame_samples = round(sample_rate * frame_ms / 1000)

    fake_vad = FakeVad([False] * 15 + [True] * 3)

    segmenter = FrameSegmenter(
        sample_rate=sample_rate,
        frame_ms=frame_ms,
        vad=fake_vad,
        pre_roll_ms=400,
        speech_start_ms=90,
        silence_ms=700,
        min_utterance_ms=300,
        max_utterance_s=30.0,
    )

    for _ in range(15):
        assert segmenter.push(make_silent_frame(frame_samples)) is None
    for _ in range(3):
        assert segmenter.push(make_speech_frame(frame_samples)) is None

    utt = None
    for _ in range(40):
        utt = segmenter.push(make_silent_frame(frame_samples))
        if utt is not None or not segmenter.speaking:
            break

    assert utt is None, "90ms blip must be dropped even though pre-roll inflates the segment"


def test_real_speech_kept_with_realistic_pre_roll():
    """Complement to the blip test: a real utterance with enough speech is kept."""
    sample_rate = 16000
    frame_ms = 30
    frame_samples = round(sample_rate * frame_ms / 1000)

    fake_vad = FakeVad([False] * 15 + [True] * 12)

    segmenter = FrameSegmenter(
        sample_rate=sample_rate,
        frame_ms=frame_ms,
        vad=fake_vad,
        pre_roll_ms=400,
        speech_start_ms=90,
        silence_ms=700,
        min_utterance_ms=300,
        max_utterance_s=30.0,
    )

    for _ in range(15):
        assert segmenter.push(make_silent_frame(frame_samples)) is None
    for _ in range(12):
        segmenter.push(make_speech_frame(frame_samples))

    utt = None
    for _ in range(40):
        utt = segmenter.push(make_silent_frame(frame_samples))
        if utt is not None:
            break

    assert utt is not None
    assert len(utt.audio) >= frame_samples * 12


def _segmenter(vad, **overrides):
    params: dict = dict(
        sample_rate=16000,
        frame_ms=30,
        vad=vad,
        pre_roll_ms=60,
        speech_start_ms=30,
        silence_ms=600,
        min_utterance_ms=30,
        max_utterance_s=30.0,
        chunk_pause_ms=300,
        soft_max_s=1.0,
        rms_floor=0.0,
        trailing_trim_ms=200,
    )
    params.update(overrides)
    return FrameSegmenter(**params)


def test_long_speech_ends_chunk_at_pause_after_soft_max():
    """A long utterance must split at a pause once soft_max_s is reached."""
    frame_samples = 480
    # 35 speech frames (> soft_max of 33) then a 10-frame (300ms) pause.
    seg = _segmenter(FakeVad([True] * 35 + [False] * 10))

    for _ in range(35):
        assert seg.push(make_speech_frame(frame_samples)) is None

    utt = None
    for _ in range(10):
        utt = seg.push(make_silent_frame(frame_samples))
        if utt is not None:
            break

    assert utt is not None, "chunk should end at the pause (silence_ms is 600, so not end-of-turn)"
    # 45 frames total; trailing 10 silence trimmed to round(200/30)=7 kept -> 3 dropped
    assert len(utt.audio) == (45 - 3) * frame_samples


def test_short_pause_before_soft_max_does_not_split():
    """A pause shorter than the threshold must not split a short utterance."""
    frame_samples = 480
    # soft_max is 10s here, so 50 frames (1.5s) + a 300ms pause must NOT finalize.
    seg = _segmenter(FakeVad([True] * 50 + [False] * 10), soft_max_s=10.0)

    for _ in range(50):
        assert seg.push(make_speech_frame(frame_samples)) is None
    for _ in range(10):
        assert seg.push(make_silent_frame(frame_samples)) is None

    # Still speaking; flush emits the whole thing.
    assert seg.flush() is not None


def test_quiet_segment_discarded_by_rms_floor():
    """Noise that trips the VAD but is too quiet should be dropped."""
    frame_samples = 480

    # High floor: the quiet utterance is dropped at finalization.
    strict = _segmenter(FakeVad([True] * 5 + [False] * 2), silence_ms=60, rms_floor=0.5)
    for _ in range(5):
        strict.push(make_speech_frame(frame_samples, value=1000))
    strict.push(make_silent_frame(frame_samples))
    assert strict.push(make_silent_frame(frame_samples)) is None

    # Same audio with a permissive floor is kept -> proves the floor is the cause.
    lenient = _segmenter(FakeVad([True] * 5 + [False] * 2), silence_ms=60, rms_floor=0.0)
    for _ in range(5):
        lenient.push(make_speech_frame(frame_samples, value=1000))
    lenient.push(make_silent_frame(frame_samples))
    assert lenient.push(make_silent_frame(frame_samples)) is not None


def test_trailing_silence_trimmed():
    """Trailing silence is trimmed to the configured padding."""
    frame_samples = 480
    # 4 speech frames then 10 silence; silence_ms=600 keeps it speaking until flush.
    seg = _segmenter(FakeVad([True] * 4 + [False] * 10))
    for _ in range(4):
        seg.push(make_speech_frame(frame_samples))
    for _ in range(10):
        assert seg.push(make_silent_frame(frame_samples)) is None

    utt = seg.flush()
    assert utt is not None
    # 14 frames; trailing 10 silence trimmed to 7 kept -> 3 dropped
    assert len(utt.audio) == (14 - 3) * frame_samples


def test_utterance_reports_rms_and_reason_for_each_end():
    """Each finalize path must label why it ended and report the audio level."""
    fs = 480

    # End of turn
    end = _segmenter(FakeVad([True] * 5 + [False] * 2), silence_ms=60)
    for _ in range(5):
        end.push(make_speech_frame(fs, value=1000))
    end.push(make_silent_frame(fs))
    u_end = end.push(make_silent_frame(fs))
    assert u_end is not None
    assert u_end.reason == "end_of_turn"
    assert u_end.rms > 0.0

    # Pause-aligned chunk
    chunk = _segmenter(FakeVad([True] * 35 + [False] * 10))
    for _ in range(35):
        chunk.push(make_speech_frame(fs, value=1000))
    u_chunk = None
    for _ in range(10):
        u_chunk = chunk.push(make_silent_frame(fs))
        if u_chunk is not None:
            break
    assert u_chunk is not None
    assert u_chunk.reason == "pause"

    # Hard cap
    capped = _segmenter(FakeVad([True] * 40), max_utterance_s=0.3)
    u_cap = None
    for _ in range(40):
        u_cap = capped.push(make_speech_frame(fs, value=1000))
        if u_cap is not None:
            break
    assert u_cap is not None
    assert u_cap.reason == "hard_cap"

    # Flush on shutdown
    flushed = _segmenter(FakeVad([True] * 3))
    for _ in range(3):
        flushed.push(make_speech_frame(fs, value=1000))
    u_flush = flushed.flush()
    assert u_flush is not None
    assert u_flush.reason == "flush"