atlas-iac/testing/tests/test_hermes_stt_streaming.py

"""Contract tests for Hermes' bounded, speculative PCM transcription stream."""

from __future__ import annotations

import importlib.util
import io
import struct
import sys
import threading
import time
from pathlib import Path
from types import SimpleNamespace

import pytest


ROOT = Path(__file__).resolve().parents[2]
SERVER = ROOT / "dockerfiles" / "hermes-jetson-stt-server.py"


def _load_server(monkeypatch):
    """Import the server without requiring the image's CUDA dependencies."""
    spec = importlib.util.spec_from_file_location("hermes_stt_stream_server", SERVER)
    assert spec and spec.loader
    module = importlib.util.module_from_spec(spec)
    monkeypatch.setitem(sys.modules, "cgi", SimpleNamespace())
    monkeypatch.setitem(
        sys.modules,
        "torch",
        SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False)),
    )
    monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace())
    spec.loader.exec_module(module)
    return module


class _Connection:
    """Capture server messages without constructing an HTTP handler."""

    def __init__(self) -> None:
        self.messages: list[dict] = []
        self.condition = threading.Condition()

    def send_json(self, payload: dict) -> None:
        with self.condition:
            self.messages.append(payload)
            self.condition.notify_all()

    def wait_for(self, message_type: str, timeout: float = 2.0) -> dict:
        deadline = time.monotonic() + timeout
        with self.condition:
            while time.monotonic() < deadline:
                match = next(
                    (
                        message
                        for message in self.messages
                        if message.get("type") == message_type
                    ),
                    None,
                )
                if match:
                    return match
                self.condition.wait(deadline - time.monotonic())
        raise AssertionError(f"no {message_type} response: {self.messages}")


class _Model:
    """Return a deterministic Whisper-shaped result and retain decode options."""

    def __init__(self) -> None:
        self.calls: list[tuple[str, dict]] = []

    def transcribe(self, path: str, **options: object) -> dict:
        self.calls.append((path, options))
        return {
            "text": " hello",
            "language": "en",
            "segments": [
                {
                    "text": " hello",
                    "no_speech_prob": 0.01,
                    "avg_logprob": -0.1,
                }
            ],
        }


def _start(session) -> None:
    session.start(
        {
            "type": "start",
            "turn_id": "turn-10",
            "format": "pcm_s16le",
            "sample_rate": 16_000,
            "language": "auto",
        }
    )


def _speech(samples: int = 1_600) -> bytes:
    return struct.pack(f"<{samples}h", *([12_000] * samples))


def _silence(samples: int = 12_000) -> bytes:
    return bytes(samples * 2)


def test_websocket_handshake_matches_rfc_example(monkeypatch):
    module = _load_server(monkeypatch)

    assert (
        module.WebSocketConnection.accept_value("dGhlIHNhbXBsZSBub25jZQ==")
        == "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="
    )


def test_websocket_read_exact_accumulates_partial_socket_reads(monkeypatch):
    module = _load_server(monkeypatch)
    payload = b'{"type":"resume"}'
    mask = b"mask"
    masked = bytes(value ^ mask[index % 4] for index, value in enumerate(payload))
    frame = bytes([0x81, 0x80 | len(payload)]) + mask + masked

    class ShortReader(io.BytesIO):
        def read(self, size: int = -1) -> bytes:
            return super().read(1 if size > 0 else size)

    handler = SimpleNamespace(rfile=ShortReader(frame), wfile=io.BytesIO())
    connection = module.WebSocketConnection(handler)

    assert connection.receive() == (0x1, payload)


def test_stream_speculation_is_reused_for_unchanged_commit(monkeypatch):
    module = _load_server(monkeypatch)
    connection = _Connection()
    model = _Model()
    session = module.StreamingTranscription(connection, model)
    _start(session)

    session.append(_speech())
    session.append(_silence())
    partial = connection.wait_for("partial")
    assert partial == {
        "type": "partial",
        "turn_id": "turn-10",
        "transcript": "hello",
        "language": "en",
        "speculative": True,
    }

    session.commit()
    final = connection.wait_for("final")
    assert final["transcript"] == "hello"
    assert final["language"] == "en"
    assert len(model.calls) == 1
    assert model.calls[0][1]["language"] is None
    assert model.calls[0][1]["condition_on_previous_text"] is False


def test_resume_invalidates_speculative_snapshot(monkeypatch):
    module = _load_server(monkeypatch)
    connection = _Connection()
    model = _Model()
    session = module.StreamingTranscription(connection, model)
    _start(session)

    session.append(_speech())
    session.append(_silence())
    connection.wait_for("partial")
    session.resume()
    session.append(_speech(800))
    session.commit()

    assert len(model.calls) == 2
    assert connection.wait_for("final")["turn_id"] == "turn-10"


def test_final_snapshot_preserves_resumed_audio_and_only_trims_tail(monkeypatch):
    module = _load_server(monkeypatch)
    session = module.StreamingTranscription(_Connection(), _Model())
    _start(session)
    voiced = _speech()
    session.append(voiced)
    session.append(_silence())

    with session._lock:
        trimmed, _ = session._snapshot_locked()
    tail_bytes = (
        module.STREAM_SAMPLE_RATE
        * module.STREAM_SAMPLE_WIDTH
        * module.VAD_TAIL_MS
        // 1000
    )
    assert len(trimmed) == len(voiced) + tail_bytes

    session.resume()
    quiet_resumed_speech = struct.pack("<800h", *([100] * 800))
    session.append(quiet_resumed_speech)
    with session._lock:
        resumed, _ = session._snapshot_locked()
    assert resumed.endswith(quiet_resumed_speech)


def test_stream_contract_rejects_wrong_format_and_oversize(monkeypatch):
    module = _load_server(monkeypatch)
    session = module.StreamingTranscription(_Connection(), _Model())
    with pytest.raises(module.WebSocketError, match="pcm_s16le"):
        session.start(
            {
                "turn_id": "turn-10",
                "format": "webm",
                "sample_rate": 16_000,
            }
        )

    malformed_rate = module.StreamingTranscription(_Connection(), _Model())
    with pytest.raises(module.WebSocketError, match="invalid sample rate"):
        malformed_rate.start(
            {
                "turn_id": "turn-10",
                "format": "pcm_s16le",
                "sample_rate": "sixteen kilohertz",
            }
        )

    session = module.StreamingTranscription(_Connection(), _Model())
    _start(session)
    monkeypatch.setattr(module, "MAX_STREAM_AUDIO_BYTES", 4)
    with pytest.raises(module.WebSocketError, match="too long"):
        session.append(bytes(6))


def test_stream_health_contract_is_declared_in_source():
    source = SERVER.read_text()

    assert '"path": "/v1/audio/transcriptions/stream"' in source
    assert '"format": "pcm_s16le"' in source
    assert '"server_vad": True' in source
    assert '"speculative": True' in source
    assert "MAX_STREAM_AUDIO_BYTES" in source