tests/test_encode_audio.py

from __future__ import annotations

import json

import pytest
import torch

from encode_audio import NativeTokenizerUnavailableError, encode_audio, main
from round_trip_test import main as round_trip_main


def _write_continuous_dav(path):
    torch.save(
        {
            "encoder.block.weight": torch.zeros(1),
            "mean_proj.weight": torch.zeros(1),
            "logs_proj.weight": torch.zeros(1),
            "dec_in_proj.weight": torch.zeros(1),
            "decoder.block.weight": torch.zeros(1),
        },
        path,
    )


def test_encode_fails_before_touching_missing_audio(tmp_path):
    checkpoint = tmp_path / "dav.pth"
    nonexistent_audio = tmp_path / "does-not-exist.wav"
    _write_continuous_dav(checkpoint)

    with pytest.raises(NativeTokenizerUnavailableError, match="missing RVQ/VQ quantizer weights") as caught:
        encode_audio(nonexistent_audio, dav_path=checkpoint)

    assert caught.value.report is not None
    assert not nonexistent_audio.exists()


def test_encode_cli_reports_blocked_json(tmp_path, capsys):
    checkpoint = tmp_path / "dav.pth"
    _write_continuous_dav(checkpoint)

    assert main([str(tmp_path / "missing.wav"), "--dav", str(checkpoint), "--json"]) == 2
    payload = json.loads(capsys.readouterr().out)
    assert payload["status"] == "BLOCKED"
    assert payload["audio_was_read"] is False


def test_round_trip_reports_blocked_before_wav_processing(tmp_path, capsys):
    checkpoint = tmp_path / "dav.pth"
    _write_continuous_dav(checkpoint)

    result = round_trip_main(
        [str(tmp_path / "missing.wav"), str(tmp_path / "out.wav"), "--dav", str(checkpoint), "--json"]
    )
    assert result == 2
    payload = json.loads(capsys.readouterr().out)
    assert payload["status"] == "BLOCKED"
    assert payload["wav_processing_started"] is False


def test_missing_encoder_is_reported_without_continuous_claim(tmp_path):
    checkpoint = tmp_path / "decoder-only.pth"
    torch.save(
        {
            "dec_in_proj.weight": torch.zeros(1),
            "decoder.block.weight": torch.zeros(1),
        },
        checkpoint,
    )

    with pytest.raises(NativeTokenizerUnavailableError) as caught:
        encode_audio(tmp_path / "not-opened.wav", dav_path=checkpoint)

    reason = str(caught.value)
    assert "waveform/tokenizer encoder weights" in reason
    assert "continuous Flow-VAE analysis path" not in reason