Owner/tests/test_barge_in.py

from __future__ import annotations

import time
import unittest

from owner_voice_pet.barge_in import AsyncBargeInMonitor, BargeInSpeakerGate
from owner_voice_pet.models import AudioFrame, AudioSegment
from owner_voice_pet.stt import MetadataSttProvider
from owner_voice_pet.transport import MemoryAudioTransport
from owner_voice_pet.vad import EnergyVadProvider


def speech_frame(idx: int, speaker_id: str, partial: str | None = "等一下") -> AudioFrame:
    metadata = {
        "duration_ms": 20,
        "speech": True,
        "speaker_id": speaker_id,
    }
    if partial is not None:
        metadata["partial_transcript"] = partial
    return AudioFrame(
        b"\xff\x7f" * 320,
        16000,
        1,
        idx * 20,
        idx,
        metadata,
    )


def segment_for_speaker(speaker_id: str) -> AudioSegment:
    return AudioSegment(
        b"\xff\x7f" * 640,
        16000,
        1,
        0,
        40,
        {"speaker_id": speaker_id},
    )


class BargeInTests(unittest.TestCase):
    def test_async_monitor_ignores_assistant_playback_voice(self) -> None:
        gate = BargeInSpeakerGate(
            enabled=True,
            user_similarity_threshold=0.62,
            assistant_reject_threshold=0.72,
            min_rms=0.001,
        )
        gate.remember_user_segment(segment_for_speaker("owner"))
        transport = MemoryAudioTransport([speech_frame(1, "assistant"), speech_frame(2, "assistant")])
        transport.start_input()
        monitor = AsyncBargeInMonitor(
            transport=transport,
            vad_provider=EnergyVadProvider(threshold=1),
            realtime_stt=MetadataSttProvider(),
            speaker_gate=gate,
            assistant_profile=gate.assistant_profile(segment_for_speaker("assistant")),
            echo_guard_ms=0,
            min_speech_ms=40,
            listen_interval_ms=1,
        )
        monitor.vad_provider.load()

        monitor.start()
        time.sleep(0.05)
        monitor.stop()

        self.assertFalse(monitor.interrupted)
        self.assertEqual(monitor.pending_frames(), [])

    def test_async_monitor_accepts_user_voice_and_buffers_frames(self) -> None:
        gate = BargeInSpeakerGate(
            enabled=True,
            user_similarity_threshold=0.62,
            assistant_reject_threshold=0.72,
            min_rms=0.001,
        )
        gate.remember_user_segment(segment_for_speaker("owner"))
        transport = MemoryAudioTransport([speech_frame(1, "owner"), speech_frame(2, "owner")])
        transport.start_input()
        monitor = AsyncBargeInMonitor(
            transport=transport,
            vad_provider=EnergyVadProvider(threshold=1),
            realtime_stt=MetadataSttProvider(),
            speaker_gate=gate,
            assistant_profile=gate.assistant_profile(segment_for_speaker("assistant")),
            echo_guard_ms=0,
            min_speech_ms=40,
            listen_interval_ms=1,
        )
        monitor.vad_provider.load()

        monitor.start()
        deadline = time.monotonic() + 1
        while not monitor.interrupted and time.monotonic() < deadline:
            time.sleep(0.005)
        monitor.stop()

        self.assertTrue(monitor.interrupted)
        self.assertEqual([frame.metadata["speaker_id"] for frame in monitor.pending_frames()], ["owner", "owner"])

    def test_user_profile_missing_falls_back_to_non_assistant_with_partial(self) -> None:
        gate = BargeInSpeakerGate(
            enabled=True,
            user_similarity_threshold=0.62,
            assistant_reject_threshold=0.72,
            min_rms=0.001,
        )
        transport = MemoryAudioTransport([speech_frame(1, "guest"), speech_frame(2, "guest")])
        transport.start_input()
        monitor = AsyncBargeInMonitor(
            transport=transport,
            vad_provider=EnergyVadProvider(threshold=1),
            realtime_stt=MetadataSttProvider(),
            speaker_gate=gate,
            assistant_profile=gate.assistant_profile(segment_for_speaker("assistant")),
            echo_guard_ms=0,
            min_speech_ms=40,
            listen_interval_ms=1,
        )
        monitor.vad_provider.load()

        monitor.start()
        deadline = time.monotonic() + 1
        while not monitor.interrupted and time.monotonic() < deadline:
            time.sleep(0.005)
        monitor.stop()

        self.assertTrue(monitor.interrupted)

    def test_async_monitor_interrupts_without_realtime_partial(self) -> None:
        gate = BargeInSpeakerGate(
            enabled=True,
            user_similarity_threshold=0.62,
            assistant_reject_threshold=0.72,
            min_rms=0.001,
        )
        gate.remember_user_segment(segment_for_speaker("owner"))
        transport = MemoryAudioTransport([speech_frame(1, "owner", None), speech_frame(2, "owner", None)])
        transport.start_input()
        monitor = AsyncBargeInMonitor(
            transport=transport,
            vad_provider=EnergyVadProvider(threshold=1),
            realtime_stt=None,
            speaker_gate=gate,
            assistant_profile=gate.assistant_profile(segment_for_speaker("assistant")),
            echo_guard_ms=0,
            min_speech_ms=40,
            listen_interval_ms=1,
        )
        monitor.vad_provider.load()

        monitor.start()
        deadline = time.monotonic() + 1
        while not monitor.interrupted and time.monotonic() < deadline:
            time.sleep(0.005)
        monitor.stop()

        self.assertTrue(monitor.interrupted)
        self.assertEqual(len(monitor.pending_frames()), 2)

    def test_render_reference_echo_does_not_interrupt(self) -> None:
        gate = BargeInSpeakerGate(
            enabled=True,
            user_similarity_threshold=0.62,
            assistant_reject_threshold=0.72,
            min_rms=0.001,
        )
        assistant = AudioSegment(b"\x01\x20\x02\x20" * 320, 16000, 1, 0, 40)
        echo_frame = AudioFrame(
            assistant.pcm[: 640 * 2],
            16000,
            1,
            20,
            1,
            {"duration_ms": 40, "speech": True},
        )
        transport = MemoryAudioTransport([echo_frame])
        transport.start_input()
        monitor = AsyncBargeInMonitor(
            transport=transport,
            vad_provider=EnergyVadProvider(threshold=1),
            realtime_stt=None,
            speaker_gate=gate,
            assistant_profile=gate.assistant_profile(assistant),
            assistant_reference=assistant,
            echo_guard_ms=0,
            min_speech_ms=40,
            listen_interval_ms=1,
        )
        monitor.vad_provider.load()

        monitor.start()
        time.sleep(0.05)
        monitor.stop()

        self.assertFalse(monitor.interrupted)


if __name__ == "__main__":
    unittest.main()