Bolt-Pet/tests/test_stt_stream.py

"""Streaming speech-to-text: the protocol, and the fallback that makes it safe
to switch on at all.

A fake websocket throughout — no network, no Deepgram account, no audio.
"""

import json
import sys
import time
from pathlib import Path

import numpy as np
import pytest

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

from bolt_pet.audio import mic
from bolt_pet.audio.stt_stream import StreamingTranscriber


class FakeSocket:
    """Records what was sent; replays scripted Deepgram frames."""

    def __init__(self, messages=(), fail_on_send=False):
        self.sent = []
        self.closed = False
        self.fail_on_send = fail_on_send
        self._messages = list(messages)

    def send_binary(self, data):
        if self.fail_on_send:
            raise ConnectionError("socket died")
        self.sent.append(data)

    def send(self, text):
        self.sent.append(text)

    def recv(self):
        if self._messages:
            return self._messages.pop(0)
        time.sleep(0.01)
        raise ConnectionError("closed")

    def close(self):
        self.closed = True


def _results(transcript, is_final=True):
    return json.dumps({
        "type": "Results", "is_final": is_final,
        "channel": {"alternatives": [{"transcript": transcript}]},
    })


def _frame(value=1000):
    return np.full(320, value, dtype=np.int16)


# ── the protocol ────────────────────────────────────────────────────────────

def test_frames_go_up_as_they_are_captured():
    socket = FakeSocket([_results("what's the weather")])
    session = StreamingTranscriber(socket)

    for _ in range(3):
        session.feed(_frame())
    text = session.finish()

    assert len(socket.sent) == 4          # three frames plus the close message
    assert text == "what's the weather"
    assert socket.closed


def test_only_final_results_are_kept():
    """Interim hypotheses change under you; concatenating them would produce
    "what what's what's the what's the weather"."""
    socket = FakeSocket([
        _results("what's", is_final=False),
        _results("what's the", is_final=False),
        _results("what's the weather", is_final=True),
    ])
    session = StreamingTranscriber(socket)
    assert session.finish() == "what's the weather"


def test_several_final_segments_are_joined():
    socket = FakeSocket([_results("turn the lights on"), _results("in the kitchen")])
    session = StreamingTranscriber(socket)
    assert session.finish() == "turn the lights on in the kitchen"


def test_junk_frames_are_ignored_rather_than_killing_the_reader():
    """An exception on the reader thread would silently end transcription for
    the rest of the utterance."""
    socket = FakeSocket(["not json at all", '{"type":"Metadata"}',
                         _results("still works")])
    session = StreamingTranscriber(socket)
    assert session.finish() == "still works"


def test_an_empty_frame_means_the_socket_closed():
    """websocket-client returns "" from recv() on a closed connection, so it
    ends the read loop rather than being treated as a blank transcript."""
    socket = FakeSocket([_results("heard this much"), "", _results("never arrives")])
    session = StreamingTranscriber(socket)
    assert session.finish() == "heard this much"


def test_a_socket_that_dies_mid_utterance_gives_up_quietly():
    socket = FakeSocket([], fail_on_send=True)
    session = StreamingTranscriber(socket)

    session.feed(_frame())      # must not raise — recording carries on
    assert session.finish() == ""


# ── opening: failure is an ordinary outcome ────────────────────────────────

def test_open_returns_none_when_it_cannot_connect():
    """None means "the one-shot path will do it", not an error."""
    def refuse():
        raise OSError("no network")

    assert StreamingTranscriber.open(connect=refuse) is None


def test_open_returns_a_session_when_it_can():
    session = StreamingTranscriber.open(connect=lambda: FakeSocket([_results("hi")]))
    assert session is not None
    assert session.finish() == "hi"


def test_streaming_is_off_without_the_switch_or_a_key(monkeypatch):
    from bolt_pet.audio import stt_stream

    monkeypatch.setattr(stt_stream.config, "STT_STREAMING", False)
    assert stt_stream.available() is False
    monkeypatch.setattr(stt_stream.config, "STT_STREAMING", True)
    monkeypatch.setattr(stt_stream.config, "DEEPGRAM_API_KEY", "")
    assert stt_stream.available() is False


# ── the capture hook ────────────────────────────────────────────────────────

class _Stream:
    """Loud frames, then quiet ones, so the VAD ends the utterance."""

    def __init__(self, loud=6, quiet=40):
        self.frames = ([np.full((320, 1), 3000, dtype=np.int16)] * loud
                       + [np.zeros((320, 1), dtype=np.int16)] * quiet)

    def read(self, n):
        return (self.frames.pop(0) if self.frames
                else np.zeros((320, 1), dtype=np.int16)), None


def test_recording_hands_every_speech_frame_to_the_listener():
    seen = []
    pcm = mic.record_utterance(_Stream(), on_frame=seen.append,
                               silence_end_sec=0.2, min_utterance_s=0.0)
    assert pcm is not None
    assert len(seen) >= 6            # every frame of speech was streamed


def test_a_listener_that_throws_cannot_break_the_recording():
    """The fallback is about to need this audio — a dead stream must not cost
    the recording too."""
    def explode(frame):
        raise RuntimeError("stream died")

    pcm = mic.record_utterance(_Stream(), on_frame=explode,
                               silence_end_sec=0.2, min_utterance_s=0.0)
    assert pcm is not None and len(pcm) > 0