tests/test_processor.py

import numpy as np
import pytest

from ark_asr_mlx.processor import MAX_AUDIO_SAMPLES, ArkASRProcessor


class _FeatureExtractor:
    hop_length = 160


def _processor_without_tokenizer() -> ArkASRProcessor:
    processor = object.__new__(ArkASRProcessor)
    processor.feature_extractor = _FeatureExtractor()
    processor.merge_factor = 4
    return processor


def test_audio_token_count_matches_encoder_downsampling() -> None:
    processor = _processor_without_tokenizer()
    assert processor.calculate_audio_token_count(16_000) == 12
    assert processor.calculate_audio_token_count(480_000) == 375


def test_empty_array_is_rejected() -> None:
    processor = _processor_without_tokenizer()
    with pytest.raises(ValueError, match="empty"):
        processor.load_audio(np.array([], dtype=np.float32), sample_rate=16_000)


def test_audio_over_thirty_seconds_is_rejected() -> None:
    processor = _processor_without_tokenizer()
    waveform = np.zeros(MAX_AUDIO_SAMPLES + 1, dtype=np.float32)
    with pytest.raises(ValueError, match="at most 30"):
        processor.load_audio(waveform, sample_rate=16_000)


def test_array_requires_sample_rate() -> None:
    processor = _processor_without_tokenizer()
    with pytest.raises(ValueError, match="sample_rate"):
        processor.load_audio(np.ones(10, dtype=np.float32))