tests/test_audio.py

"""Unit tests for capture, VAD gating, resampling and playback."""

import time

import numpy as np

from yrobot.audio import (
    AUDIO_STARTUP_CONFIG,
    FRAME_SAMPLES,
    Microphone,
    PlaybackEchoMatcher,
    Speaker,
    StreamResampler,
    UplinkGain,
    VoiceDetector,
    apply_audio_startup_config,
)


class FakeVad:
    def __init__(self, result=True):
        self.result = result

    def is_speech(self, pcm, rate):
        return self.result


class FakeMedia:
    def __init__(self):
        self.pushed = []
        self.cleared = 0
        self.samples = []
        self.audio = self

    def get_audio_sample(self):
        return self.samples.pop(0) if self.samples else None

    def push_audio_sample(self, data):
        self.pushed.append(np.asarray(data))

    def clear_player(self):
        self.cleared += 1


def test_duplex_audio_profile_uses_verified_sdk_config():
    class ConfigurableAudio:
        def __init__(self):
            self.calls = []

        def apply_audio_config(self, config, *, verify, write_settle_seconds):
            self.calls.append((config, verify, write_settle_seconds))
            return True

    media = FakeMedia()
    media.audio = ConfigurableAudio()

    assert apply_audio_startup_config(media, write_settle_seconds=0) is True
    assert media.audio.calls == [(AUDIO_STARTUP_CONFIG, True, 0)]


def test_duplex_audio_profile_is_best_effort_without_sdk_api():
    media = FakeMedia()
    media.audio = object()
    assert apply_audio_startup_config(media) is False


def test_microphone_reframes_arbitrary_stereo_blocks():
    media = FakeMedia()
    media.samples = [np.zeros((450, 2), np.float32), np.zeros((200, 2), np.float32)]
    mic = Microphone(media)
    assert [len(f) for f in mic.read_frames()] == [FRAME_SAMPLES]
    assert [len(f) for f in mic.read_frames()] == [FRAME_SAMPLES]  # 130 carried over


def test_voice_detector_needs_streak_and_energy():
    det = VoiceDetector(vad=FakeVad(True))
    loud = np.full(FRAME_SAMPLES, 0.1, np.float32)
    quiet = np.full(FRAME_SAMPLES, 1e-4, np.float32)
    assert det.process(quiet, 0.00) is False  # energy below floor gate
    assert det.process(loud, 0.02) is False  # streak 1
    assert det.process(loud, 0.04) is False  # streak 2
    assert det.process(loud, 0.06) is True  # confirmed at 3
    assert det.active(0.30) is True
    assert det.active(0.40) is False


def test_voice_detector_adapts_noise_floor():
    det = VoiceDetector(vad=FakeVad(True))
    hum = np.full(FRAME_SAMPLES, 0.02, np.float32)  # steady motor noise
    for i in range(400):
        det.process(hum, i * 0.02)
    assert det.process(hum, 9.0) is False  # floor swallowed the hum
    speech = np.full(FRAME_SAMPLES, 0.3, np.float32)
    for i in range(3):
        det.process(speech, 10.0 + i * 0.02)
    assert det.process(speech, 10.06) is True


def test_voice_detector_frozen_floor_keeps_barge_sensitivity():
    det = VoiceDetector(vad=FakeVad(True))
    echo = np.full(FRAME_SAMPLES, 0.05, np.float32)
    # 10 s of the robot's own monologue echo: the floor must not learn it
    for i in range(500):
        det.process(echo, i * 0.02, floor_frozen=True)
    speech = np.full(FRAME_SAMPLES, 0.05, np.float32)  # user at the same level
    voiced = False
    for i in range(3):
        voiced = det.process(speech, 11.0 + i * 0.02, floor_frozen=True)
    assert voiced is True  # without the freeze the floor would gate this out


def _speech_like_pcm(samples: int, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    excitation = rng.normal(size=samples + 4)
    colored = np.convolve(excitation, [0.03, 0.12, 0.35, 0.12, 0.03], mode="valid")
    envelope = np.repeat(
        np.clip(rng.normal(0.6, 0.25, samples // 160 + 1), 0.05, 1.0),
        160,
    )[:samples]
    return (colored * envelope * 0.2).astype(np.float32)


def test_playback_echo_matcher_finds_delayed_filtered_echo():
    matcher = PlaybackEchoMatcher()
    playout = _speech_like_pcm(32_000, seed=1)
    matcher.record(8.0, playout)

    source = playout[24_000:27_200]
    filtered = np.convolve(source, [0.2, 0.6, 0.2], mode="same")
    mic = (filtered * 0.12).astype(np.float32)
    match = matcher.match(mic, now=10.0)

    assert match.similarity > 0.9
    assert match.unexplained_db < -42.0
    assert 250.0 < match.lag_ms < 750.0


def test_playback_echo_matcher_leaves_unrelated_near_end_energy():
    matcher = PlaybackEchoMatcher()
    matcher.record(8.0, _speech_like_pcm(32_000, seed=2))
    user = _speech_like_pcm(3_200, seed=3) * 0.25

    match = matcher.match(user, now=10.0)

    assert match.similarity < 0.5
    assert match.unexplained_db > -42.0


def test_uplink_gain_boosts_quiet_speech_not_noise():
    agc = UplinkGain()
    speech = np.full(8000, 0.03, np.float32)  # −30 dB: typical XVF capture
    first = agc.process(speech)
    assert float(np.abs(first).max()) > 0.03  # boosting immediately
    frozen = agc.gain
    agc.process(np.full(8000, 0.001, np.float32))  # room noise: no update
    assert agc.gain == frozen
    for _ in range(20):
        out = agc.process(speech)
    assert abs(float(np.sqrt(np.mean(np.square(out)))) - UplinkGain.TARGET_RMS) < 0.02


def test_uplink_gain_never_amplifies_loud_speech_or_clips():
    agc = UplinkGain()
    loud = np.full(8000, 0.5, np.float32)
    for _ in range(10):
        out = agc.process(loud)
    assert agc.gain == 1.0
    assert float(np.abs(out).max()) <= 1.0


def test_uplink_gain_releases_in_playback_without_confirmed_user():
    agc = UplinkGain()
    quiet_speech = np.full(8000, 0.02, np.float32)
    for _ in range(10):
        agc.process(quiet_speech)
    boosted = agc.gain
    assert boosted > 2.0
    for _ in range(5):
        agc.process(
            quiet_speech,
            playback_active=True,
            confirmed_user_voice=False,
        )
    assert 1.0 < agc.gain < boosted


def test_resampler_ratio_and_continuity():
    rs = StreamResampler(24_000, 16_000)
    ramp = np.linspace(0.0, 1.0, 24_000, dtype=np.float32)
    out = np.concatenate([rs.process(chunk) for chunk in np.array_split(ramp, 13)])
    assert abs(len(out) - 16_000) <= 2
    assert np.all(np.diff(out[32:]) >= -1e-6)  # startup FIR ringing only; no seams


def test_resampler_filters_content_above_output_nyquist():
    t = np.arange(24_000) / 24_000
    passband = np.sin(2 * np.pi * 1000 * t).astype(np.float32)
    stopband = np.sin(2 * np.pi * 10_000 * t).astype(np.float32)
    low = StreamResampler().process(passband)[100:]
    high = StreamResampler().process(stopband)[100:]
    assert float(np.sqrt(np.mean(np.square(high)))) < 0.1 * float(np.sqrt(np.mean(np.square(low))))


def test_speaker_plays_after_boundary_and_flushes_on_interrupt():
    media = FakeMedia()
    speaker = Speaker(media)
    speaker.start()
    try:
        speaker.play(speaker.epoch, np.ones(2400, np.float32))  # 100 ms < preroll
        speaker.utterance_end()  # boundary flushes the short reply out
        deadline = time.monotonic() + 2.0
        while not media.pushed and time.monotonic() < deadline:
            time.sleep(0.01)
        assert media.pushed and len(media.pushed[0]) == 1600

        stale_epoch = speaker.epoch
        speaker.interrupt()
        speaker.play(stale_epoch, np.ones(24_000, np.float32))  # late, old turn
        deadline = time.monotonic() + 2.0
        while media.cleared == 0 and time.monotonic() < deadline:
            time.sleep(0.01)
        assert media.cleared == 1
        time.sleep(0.2)
        assert len(media.pushed) == 1  # stale audio never reached the device
    finally:
        speaker.close()
        speaker.join(timeout=2)


def test_speaker_reports_hard_interrupt_clear_completion():
    media = FakeMedia()
    speaker = Speaker(media)
    speaker.start()
    try:
        epoch = speaker.interrupt()
        assert speaker.wait_flushed(epoch, timeout=2.0)
        assert media.cleared == 1
    finally:
        speaker.close()
        speaker.join(timeout=2)