rehearshalhub/worker/tests/test_analyse_range.py

"""Tests for range analysis pipeline — Essentia mocked out."""

import uuid
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch

import numpy as np
import pytest

from worker.analyzers.base import AnalysisResult
from worker.pipeline.analyse_range import run_range_analysis, slice_audio


def test_slice_audio_correct_samples():
    sr = 44100
    audio = np.ones(sr * 10, dtype=np.float32)  # 10 seconds
    sliced = slice_audio(audio, sr, start_ms=2000, end_ms=5000)
    expected_len = int(3.0 * sr)  # 3 seconds
    assert abs(len(sliced) - expected_len) <= sr * 0.01  # within 10ms


def test_slice_audio_preserves_content():
    sr = 44100
    audio = np.arange(sr * 10, dtype=np.float32)
    sliced = slice_audio(audio, sr, start_ms=0, end_ms=1000)
    assert sliced[0] == 0.0
    assert len(sliced) == sr


@pytest.mark.asyncio
async def test_run_range_analysis_merges_results():
    """All analyzers are mocked — tests the merge + DB write logic."""
    from worker.analyzers.bpm import BPMAnalyzer
    from worker.analyzers.key import KeyAnalyzer
    from worker.analyzers.loudness import LoudnessAnalyzer

    audio = np.sin(np.linspace(0, 2 * np.pi * 5, 44100 * 5)).astype(np.float32)
    sr = 44100
    version_id = uuid.uuid4()
    annotation_id = uuid.uuid4()

    mock_session = AsyncMock()
    mock_session.get = AsyncMock(return_value=None)
    mock_session.add = MagicMock()
    mock_session.commit = AsyncMock()

    result_mock = AsyncMock()
    result_mock.scalar_one_or_none.return_value = None
    mock_session.execute.return_value = result_mock

    with (
        patch.object(
            BPMAnalyzer,
            "analyze",
            return_value=AnalysisResult("bpm", {"bpm": 120.0, "bpm_confidence": 0.9}),
        ),
        patch.object(
            KeyAnalyzer,
            "analyze",
            return_value=AnalysisResult("key", {"key": "A minor", "scale": "minor", "key_confidence": 0.8}),
        ),
        patch.object(
            LoudnessAnalyzer,
            "analyze",
            return_value=AnalysisResult(
                "loudness",
                {"avg_loudness_lufs": -18.0, "peak_loudness_dbfs": -6.0, "energy": 0.5},
            ),
        ),
    ):
        result = await run_range_analysis(
            audio=audio,
            sample_rate=sr,
            version_id=version_id,
            annotation_id=annotation_id,
            start_ms=0,
            end_ms=5000,
            session=mock_session,
        )

    assert result["bpm"] == 120.0
    assert result["key"] == "A minor"
    assert result["avg_loudness_lufs"] == -18.0
    mock_session.add.assert_called_once()


@pytest.mark.asyncio
async def test_run_range_analysis_handles_analyzer_failure():
    """If one analyzer raises, others should still run."""
    from worker.analyzers.bpm import BPMAnalyzer
    from worker.analyzers.chroma import ChromaAnalyzer

    audio = np.ones(44100, dtype=np.float32)

    mock_session = AsyncMock()
    result_mock = AsyncMock()
    result_mock.scalar_one_or_none.return_value = None
    mock_session.execute.return_value = result_mock
    mock_session.add = MagicMock()
    mock_session.commit = AsyncMock()

    with (
        patch.object(BPMAnalyzer, "analyze", side_effect=RuntimeError("Essentia not available")),
        patch.object(
            ChromaAnalyzer,
            "analyze",
            return_value=AnalysisResult("chroma", {"chroma_vector": [0.1] * 12}),
        ),
    ):
        result = await run_range_analysis(
            audio=audio,
            sample_rate=44100,
            version_id=uuid.uuid4(),
            annotation_id=uuid.uuid4(),
            start_ms=0,
            end_ms=1000,
            session=mock_session,
        )

    # bpm should be None (failed), chroma should be present
    assert result.get("bpm") is None
    assert result.get("chroma_vector") == [0.1] * 12


@pytest.mark.asyncio
async def test_run_range_analysis_empty_slice_raises():
    audio = np.array([], dtype=np.float32)
    with pytest.raises(ValueError, match="Empty audio slice"):
        await run_range_analysis(
            audio=audio,
            sample_rate=44100,
            version_id=uuid.uuid4(),
            annotation_id=uuid.uuid4(),
            start_ms=0,
            end_ms=0,
            session=AsyncMock(),
        )