music-assistant-server

1"""Tests for SonicAnalysisProvider._finalize base-class contract."""
2
3from __future__ import annotations
4
5import asyncio
6import struct
7import time
8from typing import cast
9from unittest.mock import AsyncMock, MagicMock
10
11import numpy as np
12import pytest
13from music_assistant_models.enums import ContentType, MediaType
14from music_assistant_models.media_items import AudioFormat
15
16from music_assistant.models.audio_analysis import AudioAnalysisData, AudioAnalysisError
17from music_assistant.providers.sonic_analysis import SonicAnalysisProvider, SonicSessionData
18from music_assistant.providers.sonic_analysis.clap_prompts import SCALAR_PROMPT_PAIRS
19
20# ---------------------------------------------------------------------------
21# Helpers
22# ---------------------------------------------------------------------------
23
24
25def _make_provider() -> tuple[SonicAnalysisProvider, AsyncMock, AsyncMock]:
26    """
27    Construct a SonicAnalysisProvider with mocked MA infrastructure.
28
29    Uses ``__new__`` to bypass ``__init__`` (no model downloads) and manually
30    sets the attributes the finalize path touches.
31
32    :returns: ``(provider, set_audio_analysis_mock, post_analysis_mock)``
33    """
34    set_aa_mock = AsyncMock()
35    post_analysis_mock = AsyncMock()
36
37    mass = MagicMock()
38    mass.streams.audio_analysis.set_audio_analysis = set_aa_mock
39    mass.streams.audio_analysis.get_audio_analysis_version = AsyncMock(return_value=None)
40    mass.streams.audio_analysis.record_analysis_failure = AsyncMock()
41    mass.create_task = MagicMock(side_effect=lambda coro: coro.close() or MagicMock())
42
43    manifest = MagicMock()
44    manifest.domain = "sonic_analysis"
45
46    p = SonicAnalysisProvider.__new__(SonicAnalysisProvider)
47    p.logger = MagicMock()
48    p.mass = mass
49    p.manifest = manifest
50    p._sessions = {}
51    p._finalize_tasks = set()
52    p.unloading = False
53    p._clap_model = None
54    p._clap_prompt_order = []
55    p._clap_text_embeddings = None
56    p.analysis_version = 1
57    p.post_analysis = post_analysis_mock  # type: ignore[method-assign]
58    return p, set_aa_mock, post_analysis_mock
59
60
61def _make_streamdetails(item_id: str = "track-1") -> MagicMock:
62    """Return a minimal streamdetails mock."""
63    sd = MagicMock()
64    sd.item_id = item_id
65    sd.provider = "test_provider"
66    sd.media_type = MediaType.TRACK
67    sd.duration = None
68    return sd
69
70
71def _make_audio_format(
72    sample_rate: int = 22050,
73    bit_depth: int = 16,
74    channels: int = 1,
75    content_type: ContentType = ContentType.PCM_S16LE,
76) -> AudioFormat:
77    """Return a real AudioFormat for 16-bit mono PCM."""
78    return AudioFormat(
79        content_type=content_type,
80        sample_rate=sample_rate,
81        bit_depth=bit_depth,
82        channels=channels,
83    )
84
85
86def _make_pcm_sine(
87    sample_rate: int = 22050,
88    duration_sec: float = 12.0,
89) -> bytes:
90    """Generate a simple sine-wave PCM block sufficient for one 10-second feature block."""
91    n = int(sample_rate * duration_sec)
92    t = np.linspace(0, duration_sec, n, endpoint=False)
93    wave = (np.sin(2 * np.pi * 440 * t) * 0.5 * 32767).astype(np.int16)
94    return struct.pack(f"<{n}h", *wave)
95
96
97def _make_session(
98    provider: SonicAnalysisProvider,
99    session_id: str,
100    *,
101    sample_rate: int = 22050,
102    bit_depth: int = 16,
103    channels: int = 1,
104) -> SonicSessionData:
105    """
106    Register a fresh SonicSessionData for session_id in the provider.
107
108    :param provider: The provider instance to register the session on.
109    :param session_id: The session ID to register.
110    :param sample_rate: PCM sample rate in Hz.
111    :param bit_depth: PCM bit depth.
112    :param channels: PCM channel count.
113    :returns: The created SonicSessionData.
114    """
115    af = _make_audio_format(sample_rate=sample_rate, bit_depth=bit_depth, channels=channels)
116    sd = _make_streamdetails()
117    block_bytes = sample_rate * (bit_depth // 8) * channels * 10
118    session = SonicSessionData(
119        streamdetails=sd,
120        audio_format=af,
121        block_bytes=block_bytes,
122        start_time=time.monotonic(),
123        clap_target_starts=[],
124        clap_target_buffers=[],
125        clap_target_complete=[],
126    )
127    provider._sessions[session_id] = session
128    return session
129
130
131# ---------------------------------------------------------------------------
132# Test 1: happy path — _finalize returns AudioAnalysisData, base class persists
133# ---------------------------------------------------------------------------
134
135
136@pytest.mark.asyncio
137async def test_finalize_calls_set_audio_analysis_and_post_analysis() -> None:
138    """
139    finalize() must call set_audio_analysis and post_analysis exactly once each.
140
141    Before the fix, _finalize returns None so the base class skips both.
142    After the fix, _finalize returns AudioAnalysisData and the base class
143    calls set_audio_analysis + post_analysis exactly once.
144    """
145    provider, set_aa, post_analysis = _make_provider()
146    session_id = "test-session-happy"
147
148    session = _make_session(provider, session_id)
149    af = session.audio_format
150    pcm = _make_pcm_sine(sample_rate=af.sample_rate, duration_sec=12.0)
151
152    # Feed PCM so we have real feature blocks
153    await provider.process_pcm_chunk(session_id, pcm)
154
155    # Call the BASE CLASS finalize (not _finalize directly)
156    await provider.finalize(session_id)
157
158    # set_audio_analysis called once by the base class
159    set_aa.assert_called_once()
160    call_kwargs = set_aa.call_args.kwargs
161    assert call_kwargs["item_id"] == session.streamdetails.item_id
162    assert call_kwargs["aa_provider_domain"] == provider.domain
163    analysis_arg = call_kwargs["analysis"]
164    assert isinstance(analysis_arg, AudioAnalysisData)
165    assert analysis_arg.duration is not None
166    assert analysis_arg.duration > 0
167
168    # post_analysis called once with the persisted analysis
169    post_analysis.assert_called_once()
170    _, post_analysis_arg = post_analysis.call_args.args
171    assert post_analysis_arg is analysis_arg
172
173    # Session cleaned up by base class
174    assert session_id not in provider._sessions
175
176
177# ---------------------------------------------------------------------------
178# Test 2: short-audio early-return — set_audio_analysis and post_analysis not called
179# ---------------------------------------------------------------------------
180
181
182@pytest.mark.asyncio
183async def test_finalize_skips_persist_when_no_feature_blocks() -> None:
184    """
185    finalize() must not persist or fire post_analysis when no feature blocks exist.
186
187    This covers the early-return path when accumulated.rms_frames is empty
188    (e.g. audio too short to produce a single 10-second block).
189    """
190    provider, set_aa, post_analysis = _make_provider()
191    session_id = "test-session-empty"
192
193    # Do NOT call process_pcm_chunk — accumulated.rms_frames stays empty
194    _make_session(provider, session_id)
195
196    await provider.finalize(session_id)
197
198    set_aa.assert_not_called()
199    post_analysis.assert_not_called()
200    # Session is still cleaned up by base class
201    assert session_id not in provider._sessions
202
203
204# ---------------------------------------------------------------------------
205# Test 3: unknown session — no exception, no persist, debug log emitted
206# ---------------------------------------------------------------------------
207
208
209@pytest.mark.asyncio
210async def test_finalize_unknown_session_is_silent() -> None:
211    """
212    finalize() on an unknown session_id must not raise and must log at debug level.
213
214    This verifies the guard at the top of _finalize for sessions that were
215    already cancelled or never started.
216    """
217    provider, set_aa, post_analysis = _make_provider()
218
219    await provider.finalize("nonexistent-session-id")
220
221    set_aa.assert_not_called()
222    post_analysis.assert_not_called()
223    # A debug log should have been emitted for the unknown session
224    provider.logger.debug.assert_called()  # type: ignore[attr-defined]
225    debug_calls = [str(c) for c in provider.logger.debug.call_args_list]  # type: ignore[attr-defined]
226    assert any("nonexistent-session-id" in c for c in debug_calls)
227
228
229# ---------------------------------------------------------------------------
230# Test 4: empty accumulated — _finalize raises AudioAnalysisError
231# ---------------------------------------------------------------------------
232
233
234@pytest.mark.asyncio
235async def test_finalize_raises_when_no_feature_blocks() -> None:
236    """_finalize must raise AudioAnalysisError when no feature blocks were accumulated."""
237    provider, _set_aa, _post_analysis = _make_provider()
238    session_id = "test-session-empty-raise"
239    _make_session(provider, session_id)
240
241    with pytest.raises(AudioAnalysisError, match="no usable audio"):
242        await provider._finalize(session_id)
243
244
245# ---------------------------------------------------------------------------
246# Test 5: a partially completed CLAP plan fails retryably instead of persisting
247# ---------------------------------------------------------------------------
248
249
250@pytest.mark.asyncio
251async def test_finalize_records_failure_when_clap_windows_incomplete() -> None:
252    """An incomplete CLAP plan must reach record_analysis_failure and suppress persistence."""
253    provider, set_aa, post_analysis = _make_provider()
254    session_id = "test-session-clap-incomplete"
255
256    session = _make_session(provider, session_id)
257    af = session.audio_format
258    # Both windows sit past the end of the 12s stream below, so neither is ever reached.
259    session.clap_target_starts = [100 * af.sample_rate, 120 * af.sample_rate]
260    session.clap_target_buffers = [[], []]
261    session.clap_target_complete = [False, False]
262
263    await provider.process_pcm_chunk(
264        session_id, _make_pcm_sine(sample_rate=af.sample_rate, duration_sec=12.0)
265    )
266
267    await provider.finalize(session_id)
268
269    set_aa.assert_not_called()
270    post_analysis.assert_not_called()
271    record_failure = cast("AsyncMock", provider.mass.streams.audio_analysis.record_analysis_failure)
272    record_failure.assert_called_once()
273    call_kwargs = record_failure.call_args.kwargs
274    assert call_kwargs["retry_at"] is not None
275    assert "0 of 2" in call_kwargs["reason"]
276    assert session_id not in provider._sessions
277
278
279# ---------------------------------------------------------------------------
280# Test 6: a sub-7s window completes through the finalize flush
281# ---------------------------------------------------------------------------
282
283
284@pytest.mark.asyncio
285async def test_finalize_flush_completes_short_window_and_persists() -> None:
286    """A window that never fills to 7s is flushed at finalize, so the analysis persists."""
287    provider, set_aa, _post_analysis = _make_provider()
288    provider.mass.create_task = MagicMock(side_effect=asyncio.create_task)  # type: ignore[method-assign]
289    session_id = "test-session-clap-flush"
290
291    session = _make_session(provider, session_id)
292    af = session.audio_format
293    # Starting at 8s of a 12s stream, this window tops out at 4s of audio.
294    session.clap_target_starts = [8 * af.sample_rate]
295    session.clap_target_buffers = [[]]
296    session.clap_target_complete = [False]
297
298    n_pairs = len(SCALAR_PROMPT_PAIRS)
299    flushed: list[int] = []
300
301    def _fake_inference(window_audio: np.ndarray, _source_sr: int) -> tuple[np.ndarray, np.ndarray]:
302        flushed.append(len(window_audio))
303        return (
304            np.ones(1024, dtype=np.float32),
305            np.zeros(2 * n_pairs, dtype=np.float32),
306        )
307
308    provider._single_window_inference_sync = _fake_inference  # type: ignore[method-assign,assignment]
309    provider._clap_model = MagicMock()
310    provider._clap_prompt_order = list(SCALAR_PROMPT_PAIRS.items())
311
312    await provider.process_pcm_chunk(
313        session_id, _make_pcm_sine(sample_rate=af.sample_rate, duration_sec=12.0)
314    )
315    await provider.finalize(session_id)
316
317    assert len(flushed) == 1
318    assert 0 < flushed[0] < 7 * af.sample_rate
319
320    set_aa.assert_called_once()
321    analysis_arg = set_aa.call_args.kwargs["analysis"]
322    assert analysis_arg.danceability is not None
323    assert analysis_arg.extra_data is not None
324    assert "clap_embedding" in analysis_arg.extra_data
325