test_composition_dedup.py 6.51 KB
"""作曲去重模块测试。

测试覆盖:
- Chromagram 提取
- 时间归一化输出维度
- Cosine 相似度计算
- 向量展开维度为 1536
"""

import os
import tempfile
import wave

import numpy as np
import pytest
from scipy.signal import resample

from composition_dedup.extractor import extract_chroma_feature, _normalize_audio_ffmpeg
from composition_dedup.similarity import (
    CompositionSimilarity,
    SimilarityDecision,
    DUPLICATE_THRESHOLD,
    SUSPECTED_THRESHOLD,
)


def _generate_test_wav(duration_sec: float = 1.0, sample_rate: int = 22050, frequency: float = 440.0) -> str:
    """生成测试用的 WAV 文件(正弦波)。

    Args:
        duration_sec: 持续时间(秒)。
        sample_rate: 采样率。
        frequency: 频率(Hz)。

    Returns:
        临时 WAV 文件路径。
    """
    t = np.linspace(0, duration_sec, int(sample_rate * duration_sec), endpoint=False)
    audio_data = (0.5 * np.sin(2 * np.pi * frequency * t)).astype(np.float32)

    tmp_path = tempfile.mktemp(suffix=".wav")
    with wave.open(tmp_path, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)  # 16-bit
        wf.setframerate(sample_rate)
        wf.writeframes((audio_data * 32767).astype(np.int16).tobytes())

    return tmp_path


class TestChromaExtraction:
    """Chromagram 提取测试。"""

    def test_extract_chroma_returns_1536_dim(self):
        """测试 Chromagram 提取返回 1536 维向量。"""
        wav_path = _generate_test_wav(duration_sec=2.0, frequency=440.0)
        try:
            feature = extract_chroma_feature(wav_path)
            assert isinstance(feature, np.ndarray)
            assert feature.shape == (1536,), f"期望 (1536,), 实际 {feature.shape}"
            assert feature.dtype == np.float32
        finally:
            if os.path.exists(wav_path):
                os.remove(wav_path)

    def test_extract_chroma_file_not_found(self):
        """测试不存在的音频文件抛出 FileNotFoundError。"""
        with pytest.raises(FileNotFoundError):
            extract_chroma_feature("/nonexistent/path/audio.mp3")

    def test_extract_chroma_different_frequencies(self):
        """测试不同频率的音频产生不同特征。"""
        wav_a = _generate_test_wav(duration_sec=2.0, frequency=440.0)
        wav_b = _generate_test_wav(duration_sec=2.0, frequency=880.0)
        try:
            feature_a = extract_chroma_feature(wav_a)
            feature_b = extract_chroma_feature(wav_b)
            # 不同频率的音频特征不应完全相同
            assert not np.allclose(feature_a, feature_b)
        finally:
            for path in [wav_a, wav_b]:
                if os.path.exists(path):
                    os.remove(path)

    def test_extract_chroma_same_audio_consistent(self):
        """测试同一音频多次提取结果一致。"""
        wav_path = _generate_test_wav(duration_sec=1.0, frequency=440.0)
        try:
            feature_1 = extract_chroma_feature(wav_path)
            feature_2 = extract_chroma_feature(wav_path)
            np.testing.assert_array_almost_equal(feature_1, feature_2, decimal=5)
        finally:
            if os.path.exists(wav_path):
                os.remove(wav_path)


class TestTimeNormalization:
    """时间归一化测试。"""

    def test_resample_chroma_to_128_frames(self):
        """测试 Chromagram 时间归一化到 128 帧。"""
        # 模拟不同长度的 Chromagram
        for num_frames in [100, 256, 512, 1000, 2000]:
            chroma = np.random.rand(12, num_frames).astype(np.float32)
            if chroma.shape[1] != 128:
                chroma = resample(chroma, 128, axis=1)
            assert chroma.shape == (12, 128), f"帧数归一化失败: {chroma.shape}"

    def test_flatten_to_1536(self):
        """测试展平后维度为 1536。"""
        chroma = np.random.rand(12, 128).astype(np.float32)
        feature = chroma.flatten()
        assert feature.shape[0] == 12 * 128 == 1536


class TestCosineSimilarity:
    """Cosine 相似度计算测试。"""

    def test_identical_vectors(self):
        """测试相同向量相似度为 1。"""
        vec = np.random.rand(1536).astype(np.float32)
        sim = CompositionSimilarity.cosine_similarity(vec, vec)
        assert abs(sim - 1.0) < 1e-6

    def test_orthogonal_vectors(self):
        """测试正交向量相似度接近 0。"""
        vec_a = np.zeros(1536)
        vec_a[0] = 1.0
        vec_b = np.zeros(1536)
        vec_b[1] = 1.0
        sim = CompositionSimilarity.cosine_similarity(vec_a, vec_b)
        assert abs(sim) < 1e-6

    def test_zero_vector(self):
        """测试零向量返回 0 相似度。"""
        vec_a = np.random.rand(1536).astype(np.float32)
        vec_b = np.zeros(1536)
        sim = CompositionSimilarity.cosine_similarity(vec_a, vec_b)
        assert sim == 0.0

    def test_similarity_range(self):
        """测试相似度值在 [0, 1] 范围内。"""
        vec_a = np.random.rand(1536).astype(np.float32)
        vec_b = np.random.rand(1536).astype(np.float32)
        sim = CompositionSimilarity.cosine_similarity(vec_a, vec_b)
        assert 0.0 <= sim <= 1.0

    def test_classify_duplicate(self):
        """测试重复判定。"""
        assert CompositionSimilarity.classify_similarity(0.96) == SimilarityDecision.DUPLICATE
        assert CompositionSimilarity.classify_similarity(0.95) == SimilarityDecision.DUPLICATE

    def test_classify_suspected(self):
        """测试疑似判定。"""
        assert CompositionSimilarity.classify_similarity(0.94) == SimilarityDecision.SUSPECTED
        assert CompositionSimilarity.classify_similarity(0.85) == SimilarityDecision.SUSPECTED

    def test_classify_new(self):
        """测试非重复判定。"""
        assert CompositionSimilarity.classify_similarity(0.84) == SimilarityDecision.NEW
        assert CompositionSimilarity.classify_similarity(0.5) == SimilarityDecision.NEW

    def test_compare_method(self):
        """测试 compare 方法同时返回相似度和判定。"""
        vec = np.random.rand(1536).astype(np.float32)
        sim, decision = CompositionSimilarity.compare(vec, vec)
        assert abs(sim - 1.0) < 1e-6
        assert decision == SimilarityDecision.DUPLICATE


class TestThresholds:
    """阈值常量测试。"""

    def test_threshold_order(self):
        """测试阈值顺序正确。"""
        assert DUPLICATE_THRESHOLD > SUSPECTED_THRESHOLD

    def test_threshold_values(self):
        """测试阈值符合设计值。"""
        assert DUPLICATE_THRESHOLD == 0.95
        assert SUSPECTED_THRESHOLD == 0.85