tests/test_audio_continuation_runner.py

"""Focused frozen-policy tests for the measured continuation runner."""

from __future__ import annotations

from pathlib import Path

from music3lab.editing.audio_continuation import CAPABILITY
from music3lab.editing.audio_continuation_runner import (
    BASELINE_CHECKPOINT_SHA256,
    SPECIALIST_CHECKPOINT_SHA256,
    _evaluation_gate,
    load_audio_continuation_config,
)


ROOT = Path(__file__).resolve().parents[1]


def _row(reference: float, continuity: float) -> dict[str, float]:
    return {
        "reference_audio_ruler": reference,
        "join_edge_rms_log_error": continuity,
        "rms": 0.1,
        "peak": 0.5,
    }


def test_config_pins_both_checkpoints_and_rejected_specialist_role() -> None:
    config, digest = load_audio_continuation_config(
        ROOT / "configs" / "audio-continuation-v1.yaml"
    )
    assert len(digest) == 64
    assert config.capability == CAPABILITY
    assert config.baseline_checkpoint_sha256 == BASELINE_CHECKPOINT_SHA256
    assert config.specialist_checkpoint_sha256 == SPECIALIST_CHECKPOINT_SHA256
    assert config.specialist_source_gate == (
        "REJECTED_teacher_latent_nmse_regression"
    )
    assert config.specialist_champion_eligible is False
    assert not any(
        (
            config.text_input_allowed,
            config.captured_condition_allowed,
            config.native_tokens_allowed,
            config.future_audio_allowed,
        )
    )


def test_gate_requires_strict_advantage_over_both_baselines_on_both_metrics() -> None:
    passing = {
        "both": _row(1.0, 0.1),
        "zero": _row(2.0, 0.2),
        "unrelated": _row(3.0, 0.3),
    }
    assert _evaluation_gate(passing)["passes"] is True
    for metric in ("reference_audio_ruler", "join_edge_rms_log_error"):
        failing = {name: dict(value) for name, value in passing.items()}
        failing["both"][metric] = failing["zero"][metric]
        gate = _evaluation_gate(failing)
        assert gate["passes"] is False


def test_gate_does_not_depend_on_specialist_label_or_champion_status() -> None:
    metrics = {
        "both": _row(2.0, 0.4),
        "zero": _row(1.0, 0.5),
        "unrelated": _row(3.0, 0.3),
    }
    gate = _evaluation_gate(metrics)
    assert gate == {
        "reference_beats_zero_and_unrelated": False,
        "continuity_beats_zero_and_unrelated": False,
        "passes": False,
    }