tests/test_silence_segmentation_security.py

import ast
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import pytest
import torch

from quality_runtime import (
    WhisperRuntime,
    _split_short_whisper_audio_at_sustained_pause,
    _split_whisper_audio,
    transcribe_whisper,
)


SAMPLE_RATE = 1_000
MINIMUM_SEGMENT_SAMPLES = 1_250
MAXIMUM_SEGMENT_SAMPLES = 28_000
MAXIMUM_VERIFIER_SEGMENTS = 12
MAXIMUM_ASR_MICROBATCH_SEGMENTS = 6


def test_space_does_not_emit_recognized_transcript_content_in_logs():
    app_source = (Path(__file__).parents[1] / "app.py").read_text(encoding="utf-8")

    assert "debug_transcript" not in app_source
    assert "transcript_text!r" not in app_source


def _app_literal(name: str):
    source = (Path(__file__).parents[1] / "app.py").read_text(encoding="utf-8")
    tree = ast.parse(source)
    for statement in tree.body:
        if not isinstance(statement, ast.Assign):
            continue
        if any(
            isinstance(target, ast.Name) and target.id == name
            for target in statement.targets
        ):
            return ast.literal_eval(statement.value)
    raise AssertionError(f"missing app literal {name}")


def _constant(seconds: float, amplitude: float = 0.2) -> np.ndarray:
    return np.full(round(seconds * SAMPLE_RATE), amplitude, dtype=np.float32)


def _pause(seconds: float) -> np.ndarray:
    return np.zeros(round(seconds * SAMPLE_RATE), dtype=np.float32)


def test_pause_threshold_is_inclusive_at_250_ms():
    def waveform(pause_samples: int) -> np.ndarray:
        return np.concatenate(
            (_constant(2.0), np.zeros(pause_samples, dtype=np.float32), _constant(2.0))
        )

    below = _split_short_whisper_audio_at_sustained_pause(
        waveform(249), sample_rate=SAMPLE_RATE
    )
    boundary = _split_short_whisper_audio_at_sustained_pause(
        waveform(250), sample_rate=SAMPLE_RATE
    )

    assert len(below) == 1
    assert len(boundary) == 2


def test_segmentation_is_deterministic_lossless_and_contiguous():
    waveform = np.concatenate(
        (_constant(2.0), _pause(0.25), _constant(2.0, 0.15))
    )

    first = _split_whisper_audio(waveform, sample_rate=SAMPLE_RATE)
    second = _split_whisper_audio(waveform.copy(), sample_rate=SAMPLE_RATE)

    assert [segment.size for segment in first] == [segment.size for segment in second]
    assert all(np.array_equal(left, right) for left, right in zip(first, second))
    assert np.array_equal(np.concatenate(first), waveform)


def test_pure_silence_is_not_turned_into_meaningless_asr_batch_segments():
    waveform = np.zeros(MAXIMUM_SEGMENT_SAMPLES, dtype=np.float32)

    segments = _split_whisper_audio(waveform, sample_rate=SAMPLE_RATE)

    assert len(segments) == 1


def test_peak_transient_does_not_reclassify_continuous_low_energy_speech_as_pause():
    time = np.arange(6 * SAMPLE_RATE, dtype=np.float32) / SAMPLE_RATE
    waveform = (0.01 * np.sin(2.0 * np.pi * 220.0 * time)).astype(np.float32)
    waveform[2_990:3_010] = 0.9

    segments = _split_whisper_audio(waveform, sample_rate=SAMPLE_RATE)

    assert len(segments) == 1


def test_every_nonempty_verifier_segment_respects_minimum_and_target_cap():
    waveform = np.concatenate(
        (_constant(27.65), _pause(0.30), _constant(0.20))
    )

    segments = _split_whisper_audio(waveform, sample_rate=SAMPLE_RATE)

    assert np.array_equal(np.concatenate(segments), waveform)
    assert len(segments) == 1
    assert all(
        MINIMUM_SEGMENT_SAMPLES <= segment.size <= 30_000
        for segment in segments
    )


def test_adversarial_short_pause_density_fails_closed_above_segment_cap():
    waveform = np.concatenate(
        tuple(
            part
            for _ in range(13)
            for part in (_constant(1.30), _pause(0.25))
        )
    )

    with pytest.raises(ValueError, match="pause segment cap"):
        _split_whisper_audio(waveform, sample_rate=SAMPLE_RATE)


def test_long_nine_pause_atoms_remain_independent_verified_contexts():
    expected_seconds = (5.72, 10.49, 4.87, 6.35, 2.95, 2.94, 6.43, 5.53, 10.17)
    waveform = _constant(sum(expected_seconds))
    cursor = 0
    for duration in expected_seconds[:-1]:
        cursor += round(duration * SAMPLE_RATE)
        waveform[cursor - 150 : cursor + 150] = 0.0

    segments = _split_whisper_audio(waveform, sample_rate=SAMPLE_RATE)

    assert len(segments) == len(expected_seconds)
    assert np.array_equal(np.concatenate(segments), waveform)
    assert [segment.size / SAMPLE_RATE for segment in segments] == pytest.approx(
        expected_seconds,
        abs=0.02,
    )


def test_h11_shorter_atomic_layout_preserves_every_proven_pause_atom():
    atomic_seconds = (5.715, 10.49, 4.865, 5.945, 2.945, 2.935, 6.425, 5.53, 9.655)
    waveform = _constant(sum(atomic_seconds))
    cursor = 0
    for duration in atomic_seconds[:-1]:
        cursor += round(duration * SAMPLE_RATE)
        waveform[cursor - 150 : cursor + 150] = 0.0

    segments = _split_whisper_audio(waveform, sample_rate=SAMPLE_RATE)

    assert len(segments) == len(atomic_seconds)
    assert np.array_equal(np.concatenate(segments), waveform)
    assert [segment.size / SAMPLE_RATE for segment in segments] == pytest.approx(
        atomic_seconds,
        abs=0.02,
    )


def test_h06_network_seam_from_app_contract_splits_exactly_two_segments():
    network_pause_seconds = _app_literal("NETWORK_INTERNAL_SILENCE_MS") / 1000.0
    waveform = np.concatenate(
        (_constant(4.51), _pause(network_pause_seconds), _constant(2.55))
    )

    segments = _split_whisper_audio(waveform, sample_rate=SAMPLE_RATE)

    assert network_pause_seconds == 0.40
    assert len(segments) == 2
    assert np.array_equal(np.concatenate(segments), waveform)


def test_long_continuous_voiced_audio_without_qualified_pause_fails_closed():
    waveform = _constant(61.0)

    with pytest.raises(ValueError, match="no qualified 250ms pause"):
        _split_whisper_audio(waveform, sample_rate=SAMPLE_RATE)


def test_long_coarse_split_allows_exactly_twelve_segments_and_rejects_thirteen():
    def waveform(context_count: int) -> np.ndarray:
        return np.concatenate(
            tuple(
                part
                for index in range(context_count)
                for part in (
                    _constant(26.70),
                    *((_pause(0.30),) if index + 1 < context_count else ()),
                )
            )
        )

    twelve = _split_whisper_audio(waveform(12), sample_rate=SAMPLE_RATE)

    assert len(twelve) == MAXIMUM_VERIFIER_SEGMENTS
    with pytest.raises(ValueError, match="pause segment cap"):
        _split_whisper_audio(waveform(13), sample_rate=SAMPLE_RATE)


class _RecordingProcessor:
    def __init__(self) -> None:
        self.calls: list[np.ndarray | list[np.ndarray]] = []

    def __call__(self, audio, **_kwargs):
        self.calls.append(audio)
        batch = len(audio) if isinstance(audio, list) else 1
        return SimpleNamespace(
            input_features=torch.zeros((batch, 80, 8), dtype=torch.float32),
            attention_mask=torch.ones((batch, 8), dtype=torch.long),
        )

    @staticmethod
    def batch_decode(token_ids, **_kwargs):
        return ["完整"] * int(token_ids.shape[0])


class _FakeModel:
    @staticmethod
    def generate(features, **_kwargs):
        return torch.ones((int(features.shape[0]), 2), dtype=torch.long)


def test_verification_asr_uses_bounded_microbatches_not_one_unbounded_batch():
    sample_rate = 16_000
    waveform = np.concatenate(
        tuple(
            part
            for index in range(7)
            for part in (
                np.full(round(26.70 * sample_rate), 0.2, dtype=np.float32),
                *(
                    (np.zeros(round(0.30 * sample_rate), dtype=np.float32),)
                    if index < 6
                    else ()
                ),
            )
        )
    )
    processor = _RecordingProcessor()
    runtime = WhisperRuntime(
        processor,
        _FakeModel(),
        torch.device("cpu"),
        torch.float32,
    )

    transcribe_whisper(waveform, sample_rate, runtime=runtime)

    batch_sizes = [len(audio) if isinstance(audio, list) else 1 for audio in processor.calls]
    assert batch_sizes == [6, 1]
    assert max(batch_sizes) <= MAXIMUM_ASR_MICROBATCH_SEGMENTS
    assert sum(batch_sizes) <= MAXIMUM_VERIFIER_SEGMENTS


def test_long_nine_atom_regression_uses_two_bounded_pause_atom_batches():
    sample_rate = 16_000
    expected_seconds = (5.72, 10.49, 4.87, 6.35, 2.95, 2.94, 6.43, 5.53, 10.17)
    waveform = np.full(round(sum(expected_seconds) * sample_rate), 0.2, np.float32)
    cursor = 0
    for duration in expected_seconds[:-1]:
        cursor += round(duration * sample_rate)
        waveform[
            cursor - round(0.15 * sample_rate) : cursor + round(0.15 * sample_rate)
        ] = 0.0
    processor = _RecordingProcessor()
    runtime = WhisperRuntime(
        processor,
        _FakeModel(),
        torch.device("cpu"),
        torch.float32,
    )

    transcript = transcribe_whisper(waveform, sample_rate, runtime=runtime)

    assert [
        len(audio) if isinstance(audio, list) else 1
        for audio in processor.calls
    ] == [6, 3]
    assert transcript.split() == ["完整"] * 9