Bolt-Assistant/tests/test_assistant.py

"""Bolt Assistant without a mic, speaker or server: synthetic audio frames,
a fake transcriber, fake converse/tts/playback. Covers wake matching, the
voice-activity detector, the listen -> ask -> speak flow, and the command
rules."""

import array
import json
import sys
from pathlib import Path
from unittest.mock import MagicMock, patch

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

import pytest

from bolt_assistant import config, server
from bolt_assistant.assistant import Assistant, annotate
from bolt_assistant.listener import Segmenter
from bolt_assistant.stt import Transcription, parse_result
from bolt_assistant.tools import Tools
from bolt_assistant.wake import find_wake

LOUD, QUIET = 3000, 20


def frame(level, samples=1280):
    return array.array("h", [level, -level] * (samples // 2)).tobytes()


# ── wake phrases ────────────────────────────────────────────────────────────

W = config.WAKE_WORDS


@pytest.mark.parametrize("said, expected", [
    ("Hey Bolt, what's the weather?", (True, "what's the weather?")),
    ("Bolt.", (True, "")),
    ("Thunderbolt turn on the lights", (True, "turn on the lights")),
    ("hey thunderbolt!", (True, "")),
    ("um, hey Bolt what time is it", (True, "what time is it")),
    ("Hey assistant. Set a timer", (True, "Set a timer")),
    ("Hey volt, play music", (True, "play music")),   # a common mishearing, after "hey"
    ("okay so bolt, hi", (True, "hi")),
    ("tighten that bolt", (False, "")),
    ("the bolt on the door", (False, "")),
    ("I was bold", (False, "")),                       # the mishearing alone isn't a wake word
    ("boltage is high", (False, "")),
    ("", (False, "")),
])
def test_find_wake(said, expected):
    assert find_wake(said, W) == expected


@pytest.mark.parametrize("said, expected", [
    ("assistant set a timer", (True, "assistant set a timer")),
    ("What's the weather, assistant?", (True, "What's the weather, assistant?")),
    ("can the assistant turn it up", (True, "can the assistant turn it up")),
    ("Hey assistant. Set a timer", (True, "Set a timer")),  # at the start: stripped as before
    ("Assistant.", (True, "")),                             # on its own: wait for the request
    ("um, assistant", (True, "")),
    ("my assistants are busy", (False, "")),                # whole words only
    ("tighten that bolt", (False, "")),
])
def test_find_wake_anywhere(said, expected):
    assert find_wake(said, W, config.ANYWHERE_WAKE_WORDS) == expected


# ── voice-activity detection ────────────────────────────────────────────────

def seg(**kw):
    args = dict(threshold=500, start_frames=3, silence_end_sec=0.4, min_sec=0.3, max_sec=2.0)
    args.update(kw)
    return Segmenter(**args)


def feed(s, frames):
    events = []
    for f in frames:
        events += s.feed(f)
    return events


def test_silence_never_starts_an_utterance():
    assert feed(seg(), [frame(QUIET)] * 50) == []


def test_speech_starts_with_pre_roll_and_ends_after_silence():
    s = seg()
    events = feed(s, [frame(QUIET)] * 6 + [frame(LOUD)] * 6 + [frame(QUIET)] * 5)
    kinds = [e.kind for e in events]
    assert kinds[0] == "start" and kinds[-1] == "end"
    start = events[0]
    assert len(start.frames) == 4 + 3  # pre-roll kept the first syllable, plus the trigger
    end = events[-1]
    assert len(end.audio) == (4 + 6 + 5) * 2560  # pre-roll + speech + trailing silence


def test_a_click_is_not_speech_and_a_short_burst_is_dropped():
    s = seg(min_sec=0.5)  # needs 6+ loud frames
    assert feed(s, [frame(LOUD), frame(QUIET)] * 5) == []  # isolated clicks
    events = feed(s, [frame(LOUD)] * 3 + [frame(QUIET)] * 5)
    assert [e.kind for e in events if e.kind != "frame"] == ["start", "dropped"]


def test_long_speech_is_cut_off_at_the_cap():
    events = feed(seg(max_sec=1.0), [frame(LOUD)] * 30)
    ends = [e for e in events if e.kind == "end"]
    assert ends and ends[0].seconds <= 1.0 + 1e-6


# ── the flow ────────────────────────────────────────────────────────────────

class FakeTranscription:
    """Each utterance takes the next scripted transcript. Bound at creation
    (the "start" event, on the main thread) because finish() now runs on a
    worker thread per utterance and two may be in flight at once."""

    script = []

    def __init__(self):
        self.frames = 0
        self.cancelled = False
        self.transcript = FakeTranscription.script.pop(0)

    def add(self, frames):
        self.frames += len(frames)

    def finish(self, audio):
        return self.transcript

    def cancel(self):
        self.cancelled = True


class Clock:
    def __init__(self):
        self.t = 1000.0

    def __call__(self):
        return self.t


def utterance():
    return [frame(LOUD)] * 8 + [frame(QUIET)] * 12


def make(replies, commands="read", streaming=False, transcription=FakeTranscription, **kw):
    asked, played, logs = [], [], []

    def converse(message, on_command, on_say=None):
        asked.append(message)
        reply = replies.pop(0)
        if callable(reply):
            return reply(on_command)
        return server.Reply(reply)

    clock = Clock()
    a = Assistant(transcription=transcription, converse=converse,
                  converse_stream=None if not streaming else kw.pop("stream"),
                  tts=lambda text, voice: [text.encode()], play=lambda chunks: played.extend(chunks),
                  tools=Tools(mode=commands, log=logs.append), log=logs.append, clock=clock, **kw)
    return a, asked, played, logs, clock


def say(a, transcript):
    FakeTranscription.script.append(transcript)
    for f in utterance():
        a.process(f)
    a.handle_transcripts(wait=True)  # what run() does between frames


def spoken(played):
    return [p.decode() for p in played if not p.startswith(b"\x00") and len(p) < 2000]


@pytest.fixture(autouse=True)
def _config(monkeypatch):
    monkeypatch.setattr(config, "STREAMING_REPLIES", False)
    monkeypatch.setattr(config, "VAD_RMS_THRESHOLD", 500)
    monkeypatch.setattr(config, "REQUIRE_WAKE_WORD", False)
    FakeTranscription.script = []


def test_ignores_speech_without_the_wake_phrase_when_required():
    a, asked, played, logs, _ = make([], wake_required=True)
    say(a, "so anyway I was telling him about the bolt")
    assert asked == [] and played == []


def test_speech_without_a_wake_phrase_is_a_request_by_default():
    a, asked, played, logs, _ = make(["It's three o'clock."])
    say(a, "what time is it?")
    assert asked[0].startswith("what time is it?\n\n[Bolt Assistant: a voice-only assistant")
    assert "It's three o'clock." in spoken(played)
    assert a.state == "idle"


def test_the_banner_says_which_mode_is_on():
    a, _, _, _, _ = make([])
    assert "anything you say" in a.listening_banner()
    required, _, _, _, _ = make([], wake_required=True)
    assert "Hey Bolt" in required.listening_banner()


def test_wake_phrase_and_request_in_one_breath():
    a, asked, played, logs, _ = make(["It's three o'clock."])
    say(a, "Hey Bolt, what time is it?")
    assert asked[0].startswith("what time is it?\n\n[Bolt Assistant: a voice-only assistant")
    assert "It's three o'clock." in spoken(played)
    assert a.state == "idle"


def test_bare_wake_phrase_chimes_then_takes_the_next_utterance():
    a, asked, played, logs, clock = make(["Sure."])
    say(a, "Hey Bolt")
    assert a.state == "command" and asked == []
    assert played  # the chime
    clock.t += 1
    say(a, "turn it up")  # no wake phrase needed now
    assert asked[0].startswith("turn it up")


def test_the_command_window_closes():
    a, asked, played, logs, clock = make([], wake_required=True)
    say(a, "Thunderbolt")
    clock.t += config.COMMAND_WAIT_SEC + 5
    say(a, "so then we left")
    assert asked == [] and a.state == "idle"


def test_a_question_keeps_listening_for_the_answer():
    a, asked, played, logs, clock = make(["Which room, the kitchen or the bedroom?", "Done."])
    say(a, "Hey Bolt, turn on the lights")
    assert a.state == "follow_up"
    clock.t += 2
    say(a, "the kitchen")
    assert asked[1].startswith("the kitchen") and a.state == "idle"


def test_mic_audio_during_speech_is_discarded():
    a, asked, played, logs, clock = make(["Hello there."])
    a.on_mic_frame(frame(LOUD))  # queued while... (simulating backlog)
    say(a, "hey bolt hello")
    assert a.frames.empty()  # what the mic heard while Bolt talked is gone
    assert not a.segmenter.in_speech


def test_listening_continues_while_a_transcript_is_in_flight():
    import threading

    release = threading.Event()

    class Slow(FakeTranscription):
        def finish(self, audio):
            release.wait(timeout=5)
            return self.transcript

    a, asked, played, logs, _ = make(["One.", "Two."], transcription=Slow)
    FakeTranscription.script += ["hey bolt first thing", "hey bolt second thing"]
    for f in utterance() + utterance():  # the second said while the first transcribes
        a.process(f)
    assert asked == []  # both transcripts still at the "server"
    release.set()
    a.handle_transcripts(wait=True)
    assert len(asked) == 2  # nothing said during the wait was lost
    assert asked[0].startswith("first thing") and asked[1].startswith("second thing")


def test_streamed_sentences_are_spoken_as_they_arrive():
    def stream(message, on_say, on_command):
        on_say("First.")
        on_say("Second.")
        return server.Reply("First. Second.", streamed=True)

    a, asked, played, logs, _ = make([], streaming=True, stream=stream)
    a._stream = stream
    say(a, "hey bolt tell me two things")
    assert spoken(played)[-2:] == ["First.", "Second."]


def test_server_failure_chimes_and_goes_back_to_listening():
    def boom(on_command):
        raise server.ServerError("down")

    a, asked, played, logs, _ = make([boom])
    say(a, "hey bolt anything")
    assert any("down" in line for line in logs) and a.state == "idle"


def test_markdown_is_not_read_aloud():
    a, asked, played, logs, _ = make(["**Sure** — see https://example.com for `details`."])
    say(a, "hey bolt help")
    text = spoken(played)[-1]
    assert "**" not in text and "https" not in text and "`" not in text


def test_the_note_describes_what_bolt_may_do():
    assert "can't run shell commands" in annotate("x", "read")
    assert "`command: <shell command>`" in annotate("x", "all")
    assert "can't act on this computer" in annotate("x", "off")


# ── commands relayed by the server ──────────────────────────────────────────

def filectl(**payload):
    return "filectl " + json.dumps(payload)


def test_read_mode_lists_files_but_refuses_shell_and_writes(tmp_path):
    (tmp_path / "a.txt").write_text("hi")
    ran = []
    tools = Tools(mode="read", log=lambda _m: None, shell=lambda c: ran.append(c) or "[exit 0]\n")
    assert "a.txt" in tools(filectl(op="list", path=str(tmp_path)))
    assert "hi" in tools(filectl(op="read", path=str(tmp_path / "a.txt")))
    assert "Not run" in tools("rm -rf /")
    assert "Not run" in tools(filectl(op="write", path=str(tmp_path / "b.txt"), content="x"))
    assert ran == [] and not (tmp_path / "b.txt").exists()


def test_all_mode_runs_commands_and_off_runs_nothing(tmp_path):
    ran = []
    shell = lambda c: ran.append(c) or "[exit 0]\n"  # noqa: E731
    assert Tools(mode="all", log=lambda _m: None, shell=shell)("command: echo hi") == "[exit 0]\n"
    assert ran == ["echo hi"]  # the stray "command:" label was stripped
    assert "Not run" in Tools(mode="off", log=lambda _m: None, shell=shell)(filectl(op="list", path="."))
    assert "isn't available" in Tools(mode="all", log=lambda _m: None)("petctl emote wave")


def test_unknown_mode_falls_back_to_read():
    assert Tools(mode="everything", log=lambda _m: None).mode == "read"


def test_a_relayed_command_goes_through_the_rules(tmp_path):
    (tmp_path / "notes.txt").write_text("buy milk")

    def reply_with_tool(on_command):
        output = on_command(filectl(op="read", path=str(tmp_path / "notes.txt")))
        assert "buy milk" in output
        return server.Reply("Your note says buy milk.")

    a, asked, played, logs, _ = make([reply_with_tool])
    say(a, "hey bolt read my notes")
    assert "Your note says buy milk." in spoken(played)


# ── server pieces ───────────────────────────────────────────────────────────

def test_whole_samples_carries_an_odd_byte_across_chunks():
    out = b"".join(server.whole_samples([b"\x01", b"\x02\x03", b"\x04\x05"]))
    assert out == b"\x01\x02\x03\x04"  # the lone trailing byte waits for its partner


def test_unsaid_drops_what_was_already_spoken():
    reply = server.unsaid(server.Reply("One. Two.\n\nThree."), ["One.", "Two."])
    assert reply == server.Reply("Three.", streamed=False)


def test_tts_needs_a_voice(monkeypatch):
    monkeypatch.setattr(config, "VOICE_ID", "")
    with pytest.raises(server.TtsError, match="ELEVENLABS_VOICE_ID"):
        server.tts_stream("hi")


def test_tts_request_shape(monkeypatch):
    monkeypatch.setattr(config, "SERVER_URL", "http://s")
    monkeypatch.setattr(config, "API_KEY", "k")
    monkeypatch.setattr(config, "VOICE_ID", "v1")
    response = MagicMock()
    response.iter_content.return_value = [b"\x00\x01"]
    with patch.object(server.requests, "post", return_value=response) as post:
        assert list(server.tts_stream("hello", voice_id="v2")) == [b"\x00\x01"]
    assert post.call_args.args[0] == "http://s/desk/tts"
    assert post.call_args.kwargs["json"]["voice_id"] == "v2"  # a speak_as pick wins
    assert post.call_args.kwargs["headers"] == {"X-Desk-Api-Key": "k"}


def test_parse_result():
    msg = json.dumps({"type": "Results", "is_final": True, "channel": {"alternatives": [{"transcript": "hey bolt"}]}})
    assert parse_result(msg) == ("hey bolt", True)
    assert parse_result("junk") == ("", False)


def test_transcription_falls_back_to_one_upload():
    class Sock:
        def __init__(self):
            self.audio, self.replies = [], [json.dumps({"is_final": True, "channel": {
                "alternatives": [{"transcript": "hey bolt"}]}}), ""]

        def send_binary(self, data):
            self.audio.append(data)

        def send(self, _text):
            pass

        def recv(self):
            return self.replies.pop(0)

        def close(self):
            pass

    late = Sock()
    calls = iter([ConnectionError("down"), late])

    def connect(_rate):
        item = next(calls)
        if isinstance(item, Exception):
            raise item
        return item

    t = Transcription(connect)
    t.add([frame(LOUD)])  # nowhere to go: the live socket never opened
    assert t.finish(frame(LOUD) * 2) == "hey bolt"
    assert b"".join(late.audio) == frame(LOUD) * 2


# ── heartbeat announcements ──────────────────────────────────────────────────

def test_announcements_are_spoken_when_idle():
    a, asked, played, logs, _ = make([])
    a.announce("You have a reminder.")
    a._drain_announcements()
    assert "You have a reminder." in spoken(played)
    assert a.state == "idle"


def test_announcements_wait_for_idle():
    a, asked, played, logs, _ = make(["It's three o'clock."])
    a.state = "busy"
    a.announce("You have a reminder.")
    a._drain_announcements()
    assert spoken(played) == []  # held back while Bolt is talking
    a.state = "idle"
    a._drain_announcements()
    assert "You have a reminder." in spoken(played)


def test_announcements_hold_while_a_transcript_is_in_flight():
    a, asked, played, logs, _ = make([])
    a._pending.append(object())  # an utterance is still at the server
    a.announce("You have a reminder.")
    a._drain_announcements()
    assert spoken(played) == []
    a._pending.clear()
    a._drain_announcements()
    assert "You have a reminder." in spoken(played)


def test_announcements_hold_while_someone_is_speaking():
    a, asked, played, logs, _ = make([])
    FakeTranscription.script.append("")  # the start event still mints a transcriber
    for f in [frame(LOUD)] * 3:
        a.process(f)  # the segmenter is mid-utterance
    a.announce("You have a reminder.")
    a._drain_announcements()
    assert spoken(played) == []
    a.segmenter.reset()
    a._drain_announcements()
    assert "You have a reminder." in spoken(played)


def test_blank_announcements_are_dropped():
    a, asked, played, logs, _ = make([])
    a.announce("   ")
    a._drain_announcements()
    assert spoken(played) == []