tests/test_audio_analysis.py

import numpy as np

from audio_analysis import analyze_waveform


def _click_track(bpm: float, duration_sec: float, sample_rate: int) -> np.ndarray:
    total = int(duration_sec * sample_rate)
    y = np.zeros(total, dtype=np.float32)

    beat_interval = max(1, int(round(sample_rate * 60.0 / bpm)))
    pulse_len = max(8, int(sample_rate * 0.03))
    pulse = np.hanning(pulse_len).astype(np.float32)

    for i in range(0, total - pulse_len, beat_interval):
        y[i : i + pulse_len] += pulse

    return y


def _tone_mix(freqs: list[float], duration_sec: float, sample_rate: int) -> np.ndarray:
    n = int(duration_sec * sample_rate)
    t = np.arange(n, dtype=np.float32) / sample_rate
    y = np.zeros(n, dtype=np.float32)

    for f in freqs:
        y += 0.22 * np.sin(2 * np.pi * f * t)
        y += 0.09 * np.sin(2 * np.pi * (2.0 * f) * t)

    return y


def test_analyze_waveform_detects_bpm_and_key_for_musical_signal():
    sample_rate = 22050
    duration = 24.0

    rhythmic = _click_track(bpm=120.0, duration_sec=duration, sample_rate=sample_rate)
    harmonic = _tone_mix(freqs=[261.63, 329.63, 392.0], duration_sec=duration, sample_rate=sample_rate)

    y = rhythmic + harmonic
    y /= np.max(np.abs(y)) + 1e-9

    analysis = analyze_waveform(y, sample_rate=sample_rate)

    assert analysis["bpm"] is not None
    assert abs(float(analysis["bpm"]) - 120.0) <= 3.0
    assert analysis["bpm_confidence"] is not None
    assert float(analysis["bpm_confidence"]) >= 0.22

    assert analysis["musical_key"] in {"C", "G"}
    assert analysis["key_scale"] == "major"
    assert analysis["key_confidence"] is not None
    assert float(analysis["key_confidence"]) >= 0.22


def test_analyze_waveform_returns_unknown_for_short_audio():
    sample_rate = 22050
    short = _tone_mix(freqs=[440.0], duration_sec=1.5, sample_rate=sample_rate)
    analysis = analyze_waveform(short, sample_rate=sample_rate)

    assert analysis["bpm"] is None
    assert analysis["bpm_confidence"] is None
    assert analysis["musical_key"] is None
    assert analysis["key_scale"] is None
    assert analysis["key_confidence"] is None