115 lines
3.7 KiB
Python
115 lines
3.7 KiB
Python
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()
|