"""Tests for cross-attention KV cache."""

import mlx.core as mx
import pytest

from mlx_music.models.acestep.cache import CrossAttentionCache


class TestCrossAttentionCache:
    def test_creation_empty(self):
        cache = CrossAttentionCache()
        assert len(cache) == 0
        assert not cache.has_layer(0)

    def test_store_and_retrieve(self):
        cache = CrossAttentionCache()
        k = mx.ones((1, 8, 10, 128))
        v = mx.zeros((1, 8, 10, 128))
        cache.update(layer_idx=0, key_states=k, value_states=v)

        assert cache.has_layer(0)
        assert not cache.has_layer(1)
        k_out, v_out = cache.get(0)
        assert mx.array_equal(k_out, k)
        assert mx.array_equal(v_out, v)

    def test_len_tracks_layers(self):
        cache = CrossAttentionCache()
        for i in range(24):
            cache.update(i, mx.zeros((1, 8, 10, 128)), mx.zeros((1, 8, 10, 128)))
        assert len(cache) == 24

    def test_get_missing_raises(self):
        cache = CrossAttentionCache()
        with pytest.raises(KeyError):
            cache.get(0)

    def test_reset_clears_all(self):
        cache = CrossAttentionCache()
        cache.update(0, mx.zeros((1, 8, 10, 128)), mx.zeros((1, 8, 10, 128)))
        cache.reset()
        assert len(cache) == 0
