Bolt-Pet/tests/test_mic.py

import sys
from pathlib import Path

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

import numpy as np
import pytest

from bolt_pet.audio import mic

FRAME_LEN = 320  # small for fast tests
SAMPLE_RATE = 8000


class _ScriptedStream:
    """Replays a fixed list of frames, then quiet forever."""

    def __init__(self, frames):
        self._frames = list(frames)

    def read(self, frames):
        if self._frames:
            frame = self._frames.pop(0)
        else:
            frame = np.zeros(FRAME_LEN, dtype=np.int16)
        return frame.reshape(-1, 1), False


def _loud(n=1):
    return [np.full(FRAME_LEN, 5000, dtype=np.int16) for _ in range(n)]


def _quiet(n=1):
    return [np.zeros(FRAME_LEN, dtype=np.int16) for _ in range(n)]


def test_returns_none_when_nothing_ever_gets_loud():
    stream = _ScriptedStream(_quiet(50))
    calls = {"i": 0}

    def should_continue():
        calls["i"] += 1
        return calls["i"] <= 50

    result = mic.record_utterance(
        stream, should_continue=should_continue,
        rms_threshold=300, silence_end_sec=0.5, max_utterance_s=5, min_utterance_s=0.1,
        frame_len=FRAME_LEN, sample_rate=SAMPLE_RATE,
    )
    assert result is None


def test_captures_speech_and_stops_after_trailing_silence():
    # speech, then enough silence to cross the silence_end_sec threshold
    silence_end_sec = 0.5
    silence_limit_frames = int(silence_end_sec * SAMPLE_RATE / FRAME_LEN)
    frames = _loud(5) + _quiet(silence_limit_frames + 2)
    stream = _ScriptedStream(frames)

    result = mic.record_utterance(
        stream, rms_threshold=300, silence_end_sec=silence_end_sec,
        max_utterance_s=5, min_utterance_s=0.05,
        frame_len=FRAME_LEN, sample_rate=SAMPLE_RATE,
    )
    assert result is not None
    # captured the loud frames plus the silence up to (and including) the
    # frame that crossed the silence-end threshold, but not endless silence
    assert len(result) < len(frames) * FRAME_LEN


def test_returns_none_if_utterance_shorter_than_minimum():
    frames = _loud(1) + _quiet(2)  # crosses silence limit almost immediately
    stream = _ScriptedStream(frames)
    result = mic.record_utterance(
        stream, rms_threshold=300, silence_end_sec=0.05,
        max_utterance_s=5, min_utterance_s=5.0,  # impossible to satisfy
        frame_len=FRAME_LEN, sample_rate=SAMPLE_RATE,
    )
    assert result is None


def test_stops_at_max_utterance_even_without_silence():
    max_utterance_s = 0.5
    max_frames = int(max_utterance_s * SAMPLE_RATE / FRAME_LEN)
    stream = _ScriptedStream(_loud(max_frames + 20))  # never goes quiet

    result = mic.record_utterance(
        stream, rms_threshold=300, silence_end_sec=10.0,  # would never trigger
        max_utterance_s=max_utterance_s, min_utterance_s=0.01,
        frame_len=FRAME_LEN, sample_rate=SAMPLE_RATE,
    )
    assert result is not None
    assert len(result) == max_frames * FRAME_LEN


def test_returns_none_when_should_continue_stops_before_speech():
    stream = _ScriptedStream(_quiet(100))
    result = mic.record_utterance(stream, should_continue=lambda: False,
                                   frame_len=FRAME_LEN, sample_rate=SAMPLE_RATE)
    assert result is None


def test_pcm_to_wav_bytes_round_trips_via_wave_module():
    import wave
    import io

    pcm = np.array([0, 100, -100, 32767, -32768], dtype=np.int16)
    wav_bytes = mic.pcm_to_wav_bytes(pcm, sample_rate=16000)
    with wave.open(io.BytesIO(wav_bytes), "rb") as wf:
        assert wf.getnchannels() == 1
        assert wf.getsampwidth() == 2
        assert wf.getframerate() == 16000
        frames = wf.readframes(wf.getnframes())
    assert np.frombuffer(frames, dtype=np.int16).tolist() == pcm.tolist()


# ── flushing buffered audio ─────────────────────────────────────────────────

class _BufferedStream:
    """A stream with a backlog, like PortAudio's ring buffer after the reader
    was blocked on a network call for a while."""

    def __init__(self, available):
        self.read_available = available
        self.reads = []

    def read(self, frames):
        self.reads.append(frames)
        self.read_available = max(0, self.read_available - frames)
        return np.zeros((frames, 1), dtype=np.int16), False


def test_flush_drops_exactly_what_was_buffered():
    stream = _BufferedStream(4096)
    assert mic.flush(stream) == 4096
    assert stream.reads == [4096]
    assert stream.read_available == 0


def test_flush_is_bounded_so_it_cannot_chase_a_live_stream():
    """A stream filling as fast as it drains must not spin forever."""
    stream = _BufferedStream(10 ** 9)
    dropped = mic.flush(stream, max_seconds=1.0, sample_rate=16000)
    assert dropped == 16000


def test_flush_is_a_noop_on_an_empty_or_fake_stream():
    stream = _BufferedStream(0)
    assert mic.flush(stream) == 0
    assert stream.reads == []
    assert mic.flush(_ScriptedStream([])) == 0   # no read_available at all
    assert mic.flush(None) == 0


def test_flush_swallows_a_device_error():
    class _Broken:
        read_available = 1024

        def read(self, frames):
            raise RuntimeError("device disappeared")

    assert mic.flush(_Broken()) == 0