"""Tests for debug hooks that capture intermediate tensors."""

import mlx.core as mx
import pytest

from mlx_music.debug.hooks import IntermediateCapture


class TestIntermediateCapture:
    def test_captures_tensors(self):
        cap = IntermediateCapture()
        cap.record("layer_0_self_attn", mx.ones((1, 10, 256)))
        cap.record("layer_0_cross_attn", mx.zeros((1, 10, 256)))

        assert len(cap) == 2
        assert "layer_0_self_attn" in cap
        assert cap["layer_0_self_attn"].shape == (1, 10, 256)

    def test_save_and_load(self, tmp_path):
        cap = IntermediateCapture()
        cap.record("test", mx.ones((2, 3)))

        path = tmp_path / "intermediates.npz"
        cap.save(str(path))

        cap2 = IntermediateCapture.load(str(path))
        assert "test" in cap2
        assert mx.allclose(cap2["test"], mx.ones((2, 3)))

    def test_disabled_is_noop(self):
        cap = IntermediateCapture(enabled=False)
        cap.record("test", mx.ones((2, 3)))
        assert len(cap) == 0
