tests/test_generated_capture.py

from __future__ import annotations

import hashlib
import json
import os
import sys
from pathlib import Path
from types import ModuleType, SimpleNamespace

import numpy as np
import pytest

import capture_generated_tokens as capture_module
from capture_generated_tokens import DepthCodeCapture, build_capture_metadata
from encode_audio import NativeTokenizerUnavailableError, encode_audio
from native_token_compatibility import validate_native_tokens


def _paired_row(c0: int) -> np.ndarray:
    row = np.array([c0, 1, 2, 3, 4, 5, 6, 7], dtype=np.int64)
    return np.stack((row, row), axis=0)


def test_capture_preserves_return_and_skips_exactly_one_priming_call():
    rows = iter((_paired_row(10), _paired_row(11), _paired_row(12)))
    sentinel = object()

    def official_function():
        return next(rows), sentinel

    capture = DepthCodeCapture(priming_calls=1)
    wrapped = capture.wrap(official_function)
    assert wrapped()[1] is sentinel
    assert wrapped()[1] is sentinel
    assert wrapped()[1] is sentinel

    tokens = capture.emitted_tokens()
    assert capture.captured_calls == 3
    assert tokens.shape == (2, 8)
    assert tokens[:, 0].tolist() == [11, 12]


def test_patch_context_restores_official_function(monkeypatch):
    rows = iter((_paired_row(10), _paired_row(11)))

    def official_function():
        return next(rows), np.zeros(1)

    fake_encoders = SimpleNamespace(_generate_depth_codes=official_function)
    monkeypatch.setattr(
        capture_module.importlib,
        "import_module",
        lambda _: fake_encoders,
    )

    with capture_module.capture_official_depth_codes() as capture:
        assert fake_encoders._generate_depth_codes is not official_function
        fake_encoders._generate_depth_codes()
        fake_encoders._generate_depth_codes()

    assert fake_encoders._generate_depth_codes is official_function
    assert capture.emitted_tokens().shape == (1, 8)


def test_run_capture_builds_pipeline_before_frame_rate_validation(monkeypatch, tmp_path):
    events = []
    rows = iter((_paired_row(10), _paired_row(11)))

    def official_depth_codes():
        return next(rows), np.zeros(1)

    fake_encoders = SimpleNamespace(_generate_depth_codes=official_depth_codes)

    class FakeGenerator:
        def __init__(self, device):
            events.append(("generator", device))

        def manual_seed(self, seed):
            events.append(("seed", seed))
            return self

    class FakePipe:
        sampling_rate = 44_100

        def load_components(self, *, dtype):
            events.append(("load_components", dtype))

        @property
        def frame_rate(self):
            events.append(("frame_rate", capture_module.FRAME_RATE_HZ))
            return capture_module.FRAME_RATE_HZ

        def to(self, device):
            events.append(("to", device))

        def __call__(self, **kwargs):
            events.append(("generate", kwargs["audio_duration"]))
            fake_encoders._generate_depth_codes()
            fake_encoders._generate_depth_codes()
            return np.zeros((1, 2, 160), dtype=np.float32)

    class FakeModularPipeline:
        @classmethod
        def from_pretrained(cls, model):
            events.append(("from_pretrained", model))
            return FakePipe()

    fake_torch = ModuleType("torch")
    fake_torch.bfloat16 = object()
    fake_torch.Generator = FakeGenerator

    fake_diffusers = ModuleType("diffusers")
    fake_diffusers.__file__ = str(tmp_path / "diffusers" / "__init__.py")
    fake_diffusers.ModularPipeline = FakeModularPipeline

    fake_soundfile = ModuleType("soundfile")

    def fake_write(path, audio, rate, *, format):
        Path(path).write_bytes(b"RIFF-fake-wave")
        events.append(("write_wav", Path(path), audio.shape, rate, format))

    fake_soundfile.write = fake_write

    monkeypatch.setitem(sys.modules, "torch", fake_torch)
    monkeypatch.setitem(sys.modules, "diffusers", fake_diffusers)
    monkeypatch.setitem(sys.modules, "soundfile", fake_soundfile)
    monkeypatch.setattr(
        capture_module,
        "_verified_diffusers_revision",
        lambda _module, expected: {
            "commit": expected,
            "source_kind": "git_checkout",
            "source_location": "/source",
            "tracked_clean": True,
            "relevant_source_path": "/source/encoders.py",
            "relevant_source_sha256": "a" * 64,
        },
    )
    monkeypatch.setattr(
        capture_module,
        "_resolve_model_snapshot",
        lambda *_args, **_kwargs: (
            tmp_path / "snapshots" / "resolved-model-commit",
            "resolved-model-commit",
        ),
    )
    monkeypatch.setattr(
        capture_module.importlib,
        "import_module",
        lambda _: fake_encoders,
    )
    real_replace = os.replace

    def tracked_replace(old, new):
        events.append(("replace", Path(new).name))
        real_replace(old, new)

    monkeypatch.setattr(capture_module.os, "replace", tracked_replace)

    args = SimpleNamespace(
        model="model-id",
        model_revision="requested-model-revision",
        diffusers_revision="diffusers-revision",
        local_files_only=True,
        dtype="bfloat16",
        device="cuda",
        seed=7,
        prompt="prompt",
        lyrics="[instrumental]",
        audio_duration=1.0,
        num_inference_steps=2,
        output_tokens=tmp_path / "tokens.npy",
        output_wav=tmp_path / "audio.wav",
        output_metadata=tmp_path / "metadata.json",
    )
    metadata = capture_module.run_capture(args)

    event_names = [event[0] for event in events]
    assert event_names.index("from_pretrained") < event_names.index("load_components")
    assert event_names.index("load_components") < event_names.index("frame_rate")
    assert event_names.index("frame_rate") < event_names.index("to")
    assert event_names.index("to") < event_names.index("generate")
    replaced = [event[1] for event in events if event[0] == "replace"]
    assert replaced == ["tokens.npy", "audio.wav", "metadata.json"]
    assert metadata["token_shape_frames_first"] == [1, 8]
    assert metadata["resolved_model_commit"] == "resolved-model-commit"
    assert metadata["device"] == "cuda"
    assert metadata["dtype"] == "bfloat16"
    assert np.load(args.output_tokens, allow_pickle=False).shape == (1, 8)
    saved = json.loads(args.output_metadata.read_text())
    assert saved["wav_reencoding_performed"] is False
    assert saved["artifact_sha256"]["wav"] == hashlib.sha256(
        args.output_wav.read_bytes()
    ).hexdigest()
    assert saved["artifact_sha256"]["tokens_npy"] == hashlib.sha256(
        args.output_tokens.read_bytes()
    ).hexdigest()

def test_capture_rejects_unpaired_cfg_rows():
    capture = DepthCodeCapture()
    unpaired = _paired_row(10)
    unpaired[1, 3] += 1
    wrapped = capture.wrap(lambda: (unpaired, np.zeros(1)))

    with pytest.raises(RuntimeError, match="not identical"):
        wrapped()



def test_diffusers_verifier_rejects_dirty_tracked_checkout(monkeypatch, tmp_path):
    repository = tmp_path / "diffusers-repository"
    (repository / ".git").mkdir(parents=True)
    package = repository / "src" / "diffusers"
    package.mkdir(parents=True)
    fake_diffusers = SimpleNamespace(__file__=str(package / "__init__.py"))
    outputs = iter(
        (
            SimpleNamespace(stdout="a" * 40 + "\n"),
            SimpleNamespace(stdout=" M src/diffusers/__init__.py\n"),
        )
    )
    monkeypatch.setattr(
        capture_module.subprocess,
        "run",
        lambda *args, **kwargs: next(outputs),
    )

    with pytest.raises(RuntimeError, match="dirty tracked files"):
        capture_module._verified_diffusers_revision(fake_diffusers, "a" * 40)



def _metadata_arguments() -> dict:
    return {
        "model_id": "model-id",
        "requested_model_revision": "requested-model-revision",
        "resolved_model_commit": "resolved-model-commit",
        "diffusers_identity": {
            "commit": "diffusers-revision",
            "source_kind": "git_checkout",
            "source_location": "/source",
            "tracked_clean": True,
            "relevant_source_path": "/source/encoders.py",
            "relevant_source_sha256": "a" * 64,
        },
        "device": "cuda",
        "dtype": "bfloat16",
        "prompt": "prompt",
        "lyrics": "[instrumental]",
        "requested_audio_duration_seconds": 1.0,
        "num_inference_steps": 30,
        "seed": 7,
        "sampling_rate": 32_000,
        "wav_sha256": "b" * 64,
        "tokens_npy_sha256": "c" * 64,
        "capture_script_sha256": "d" * 64,
    }


def test_metadata_proves_internal_generation_not_wav_reencoding():
    tokens = np.stack((_paired_row(11)[0], _paired_row(12)[0]))
    audio = np.zeros((2, 3_200), dtype=np.float32)
    metadata = build_capture_metadata(
        tokens=tokens,
        audio=audio,
        captured_calls=3,
        priming_rows_skipped=1,
        **_metadata_arguments(),
    )

    assert metadata["capture_kind"] == "official_internal_generation_tokens"
    assert metadata["wav_reencoding_performed"] is False
    assert metadata["model_id"] == "model-id"
    assert metadata["resolved_model_commit"] == "resolved-model-commit"
    assert metadata["diffusers_source_identity"]["tracked_clean"] is True
    assert metadata["token_shape_frames_first"] == [2, 8]
    assert metadata["audio_shape_channels_first"] == [2, 3_200]
    assert metadata["audio_duration_seconds"] == pytest.approx(0.1)


def test_metadata_rejects_call_to_frame_misalignment():
    with pytest.raises(ValueError, match="alignment mismatch"):
        build_capture_metadata(
            tokens=np.stack((_paired_row(11)[0], _paired_row(12)[0])),
            audio=np.zeros((2, 3_200), dtype=np.float32),
            captured_calls=2,
            priming_rows_skipped=1,
            **_metadata_arguments(),
        )

ARTIFACT_ENV_NAMES = (
    "MINIMAX_DAV_PATH",
    "MINIMAX_GENERATED_WAV_PATH",
    "MINIMAX_INTERNAL_TOKENS_PATH",
    "MINIMAX_GENERATION_CAPTURE_PATH",
)


def _artifact_paths() -> tuple[Path, Path, Path, Path] | None:
    values = [os.environ.get(name) for name in ARTIFACT_ENV_NAMES]
    configured = [bool(value) for value in values]
    if any(configured) and not all(configured):
        missing = [
            name for name, is_configured in zip(ARTIFACT_ENV_NAMES, configured)
            if not is_configured
        ]
        raise RuntimeError(
            "generated-capture integration requires all four artifact "
            f"environment variables; missing: {', '.join(missing)}"
        )
    if not any(configured):
        return None
    return tuple(Path(value) for value in values)


def test_partial_artifact_environment_is_rejected(monkeypatch):
    for name in ARTIFACT_ENV_NAMES:
        monkeypatch.delenv(name, raising=False)
    monkeypatch.setenv("MINIMAX_DAV_PATH", "/only/dav.pth")

    with pytest.raises(RuntimeError, match="requires all four"):
        _artifact_paths()

def test_real_generated_sample_has_valid_internal_tokens_but_wav_encoding_is_blocked():
    paths = _artifact_paths()
    if paths is None:
        pytest.skip("set all four generated-capture artifact environment variables")
    dav_path, wav_path, tokens_path, metadata_path = paths

    import soundfile as sf

    tokens = np.load(tokens_path, allow_pickle=False)
    validated = validate_native_tokens(tokens, layout="frames_first")
    metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
    audio_info = sf.info(wav_path)

    required = {
        "capture_kind",
        "wav_reencoding_performed",
        "model_id",
        "requested_model_revision",
        "resolved_model_commit",
        "diffusers_revision",
        "diffusers_source_identity",
        "capture_script_sha256",
        "device",
        "dtype",
        "prompt",
        "lyrics",
        "requested_audio_duration_seconds",
        "num_inference_steps",
        "seed",
        "artifact_sha256",
    }
    assert required <= metadata.keys()
    assert metadata["capture_kind"] == "official_internal_generation_tokens"
    assert metadata["wav_reencoding_performed"] is False
    assert metadata["model_id"] == "MiniMaxAI/MiniMax-Music3"
    assert metadata["requested_model_revision"] == (
        "fbdf52fbaaca799592917417eb05f1899f1255ec"
    )
    assert metadata["resolved_model_commit"] == (
        "fbdf52fbaaca799592917417eb05f1899f1255ec"
    )
    assert metadata["diffusers_revision"] == (
        "90b4e34e79a86ec5e7f2437634fe95ecd2108796"
    )
    source_identity = metadata["diffusers_source_identity"]
    assert source_identity["commit"] == metadata["diffusers_revision"]
    assert source_identity["source_kind"] in {"git_checkout", "vcs_install"}
    if source_identity["source_kind"] == "git_checkout":
        assert source_identity["tracked_clean"] is True
    relevant_source = Path(source_identity["relevant_source_path"])
    assert source_identity["relevant_source_sha256"] == hashlib.sha256(
        relevant_source.read_bytes()
    ).hexdigest()
    assert metadata["capture_script_sha256"] == hashlib.sha256(
        Path(capture_module.__file__).read_bytes()
    ).hexdigest()
    assert metadata["device"] == "cuda"
    assert metadata["dtype"] == "bfloat16"
    assert metadata["prompt"] == (
        "Instrumental French house, 126 BPM, E minor, filtered disco loop, "
        "punchy kick and warm bass."
    )
    assert metadata["lyrics"] == "[instrumental]"
    assert metadata["requested_audio_duration_seconds"] == 1.0
    assert metadata["num_inference_steps"] == 30
    assert metadata["seed"] == 7
    assert metadata["artifact_sha256"]["wav"] == hashlib.sha256(
        wav_path.read_bytes()
    ).hexdigest()
    assert metadata["artifact_sha256"]["tokens_npy"] == hashlib.sha256(
        tokens_path.read_bytes()
    ).hexdigest()

    assert validated.min(axis=0).tolist() == [1012, 95, 25, 42, 83, 2, 3, 67]
    assert validated.max(axis=0).tolist() == [
        16163, 984, 1005, 950, 941, 967, 984, 957
    ]
    assert validated.shape == (25, 8)
    assert metadata["captured_calls_including_priming"] == 26
    assert metadata["priming_rows_skipped"] == 1
    assert metadata["token_shape_frames_first"] == [25, 8]
    assert metadata["token_mins"] == validated.min(axis=0).tolist()
    assert metadata["token_maxs"] == validated.max(axis=0).tolist()
    assert audio_info.samplerate == 44_100
    assert audio_info.channels == 2
    assert metadata["sampling_rate"] == 44_100
    assert metadata["audio_shape_channels_first"] == [2, 44_032]
    assert metadata["audio_duration_seconds"] == pytest.approx(0.9984580498866213)
    assert audio_info.frames == 44_032

    # This is deliberately not a token self-comparison. The generated WAV is
    # passed to the released encoder API and remains blocked before WAV I/O.
    with pytest.raises(NativeTokenizerUnavailableError):
        encode_audio(wav_path, dav_path=dav_path)