STT-Assistant/tests/test_bus.py

"""Tests for the publish-only event bus and subscriber."""

import os
import shutil
import socket
import tempfile
import threading
import time

import pytest

from stt_daemon import events
from stt_daemon.bus import EventPublisher, EventSubscriber


@pytest.fixture
def sock_path():
    # Short path: AF_UNIX paths are limited to ~108 bytes.
    directory = tempfile.mkdtemp(prefix="sttb-")
    try:
        yield os.path.join(directory, "s.sock")
    finally:
        shutil.rmtree(directory, ignore_errors=True)


def _wait(predicate, timeout=3.0):
    deadline = time.time() + timeout
    while time.time() < deadline:
        if predicate():
            return True
        time.sleep(0.01)
    return predicate()


class _Collector:
    def __init__(self, subscriber):
        self.events = []
        self._thread = threading.Thread(target=self._run, args=(subscriber,), daemon=True)
        self._thread.start()

    def _run(self, subscriber):
        try:
            for event in subscriber.events():
                self.events.append(event)
        except Exception:  # noqa: BLE001 - collector must never crash the test thread
            pass

    def types(self):
        return [e["type"] for e in self.events]

    def wait_for(self, predicate, timeout=3.0):
        return _wait(lambda: predicate(self.events), timeout)


def _has(event_list, type_):
    return any(e["type"] == type_ for e in event_list)


def _count(event_list, type_):
    return sum(1 for e in event_list if e["type"] == type_)


def test_publish_subscribe_in_order(sock_path):
    pub = EventPublisher(sock_path, ring_size=10)
    pub.start()
    sub = EventSubscriber(sock_path, reconnect_delay=0.05)
    collector = _Collector(sub)

    assert collector.wait_for(lambda e: _has(e, events.STATUS))

    for _ in range(5):
        pub.emit(events.SPEECH_START)

    assert collector.wait_for(lambda e: _count(e, events.SPEECH_START) == 5)
    seqs = [e["seq"] for e in collector.events if e["type"] == events.SPEECH_START]
    assert seqs == sorted(seqs)

    sub.close()
    pub.close()


def test_emit_assigns_monotonic_seq(sock_path):
    pub = EventPublisher(sock_path, ring_size=10)
    pub.start()
    sub = EventSubscriber(sock_path, reconnect_delay=0.05)
    collector = _Collector(sub)
    assert collector.wait_for(lambda e: _has(e, events.STATUS))

    for _ in range(3):
        pub.emit(events.HEARTBEAT)
    assert collector.wait_for(lambda e: _count(e, events.HEARTBEAT) == 3)

    seqs = [e["seq"] for e in collector.events if e["type"] == events.HEARTBEAT]
    assert seqs == sorted(seqs)
    assert len(set(seqs)) == 3

    sub.close()
    pub.close()


def test_ring_replay_on_attach(sock_path):
    pub = EventPublisher(sock_path, ring_size=10)
    pub.start()

    for i in range(3):
        pub.emit(events.TRANSCRIPT, text=f"t{i}")
    assert _wait(lambda: pub.ring_len() == 3)

    sub = EventSubscriber(sock_path, reconnect_delay=0.05)
    collector = _Collector(sub)
    assert collector.wait_for(lambda e: _has(e, events.RING_END))

    types = collector.types()
    assert types[0] == events.STATUS
    ring_end_index = types.index(events.RING_END)
    assert types[1:ring_end_index] == [events.TRANSCRIPT] * 3

    sub.close()
    pub.close()


def test_two_subscribers_each_get_their_own_replay(sock_path):
    pub = EventPublisher(sock_path, ring_size=10)
    pub.start()
    pub.emit(events.TRANSCRIPT, text="a")
    assert _wait(lambda: pub.ring_len() == 1)

    subs = [EventSubscriber(sock_path, reconnect_delay=0.05) for _ in range(2)]
    collectors = [_Collector(s) for s in subs]

    for collector in collectors:
        assert collector.wait_for(lambda e: _has(e, events.RING_END))
        types = collector.types()
        assert types[0] == events.STATUS
        assert types[1:types.index(events.RING_END)] == [events.TRANSCRIPT]

    for sub in subs:
        sub.close()
    pub.close()


def test_publisher_never_blocks_on_slow_subscriber(sock_path):
    pub = EventPublisher(sock_path, ring_size=50)
    pub.start()

    # A raw client that connects and never reads: its writer blocks in sendall.
    slow = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    slow.connect(sock_path)
    assert _wait(lambda: pub.client_count() == 1)

    total = 20000
    start = time.perf_counter()
    for _ in range(total):
        pub.emit(events.HEARTBEAT, rms=0.01)
    elapsed = time.perf_counter() - start

    assert elapsed < 2.0, f"publish path blocked: {elapsed:.3f}s for {total} events"

    slow.close()
    assert _wait(lambda: pub.client_count() == 0)

    # The publisher is still healthy: a fresh subscriber attaches promptly.
    sub = EventSubscriber(sock_path, reconnect_delay=0.05)
    collector = _Collector(sub)
    assert collector.wait_for(lambda e: _has(e, events.RING_END))

    sub.close()
    pub.close()


def test_subscriber_reconnects_after_restart(sock_path):
    pub = EventPublisher(sock_path, ring_size=10)
    pub.start()
    sub = EventSubscriber(sock_path, reconnect_delay=0.05)
    collector = _Collector(sub)
    assert collector.wait_for(lambda e: _has(e, events.STATUS))

    pub.emit(events.TRANSCRIPT, text="one")
    assert collector.wait_for(lambda e: _count(e, events.TRANSCRIPT) == 1)

    pub.close()

    pub2 = EventPublisher(sock_path, ring_size=10)
    pub2.start()
    assert collector.wait_for(lambda e: _count(e, events.STATUS) >= 2, timeout=5.0)

    pub2.emit(events.TRANSCRIPT, text="two")
    assert collector.wait_for(lambda e: _count(e, events.TRANSCRIPT) == 2, timeout=5.0)

    sub.close()
    pub2.close()


def test_socket_is_private_and_removed_on_close(sock_path):
    pub = EventPublisher(sock_path, ring_size=10)
    pub.start()

    assert os.path.exists(sock_path)
    assert os.stat(sock_path).st_mode & 0o777 == 0o600

    pub.close()
    assert not os.path.exists(sock_path)


def test_status_carries_provider_fields(sock_path):
    pub = EventPublisher(
        sock_path,
        ring_size=10,
        status_provider=lambda: {"capture": "VirtualMicSink.monitor", "segments_done": 7},
    )
    pub.start()

    sub = EventSubscriber(sock_path, reconnect_delay=0.05)
    collector = _Collector(sub)
    assert collector.wait_for(lambda e: _has(e, events.STATUS))

    status = next(e for e in collector.events if e["type"] == events.STATUS)
    assert status["capture"] == "VirtualMicSink.monitor"
    assert status["segments_done"] == 7
    assert "ring_seq" in status

    sub.close()
    pub.close()