tests/test_chromagram.py

import numpy as np
from app.services.chromagram import extract_chroma_mean

def sine(freq, duration=2.0, sr=22050):
    t = np.linspace(0, duration, int(sr * duration))
    return np.sin(2 * np.pi * freq * t).astype(np.float32)

def test_output_shape():
    assert extract_chroma_mean(sine(440.0), sr=22050).shape == (12,)

def test_output_is_normalized():
    chroma = extract_chroma_mean(sine(440.0), sr=22050)
    assert abs(chroma.sum() - 1.0) < 0.05

def test_output_non_negative():
    assert (extract_chroma_mean(sine(440.0), sr=22050) >= 0).all()

def test_a440_peak_at_index_9():
    # A = pitch class 9 (C=0 … A=9)
    chroma = extract_chroma_mean(sine(440.0, duration=4.0), sr=22050)
    assert np.argmax(chroma) == 9