/
/
/
1"""Tests for the live-CLAP finalize + cancel paths."""
2
3from __future__ import annotations
4
5import asyncio
6import math
7from unittest.mock import MagicMock
8
9import numpy as np
10import pytest
11
12from music_assistant.helpers.datetime import utc
13from music_assistant.models.audio_analysis import AudioAnalysisData, AudioAnalysisError
14from music_assistant.providers.sonic_analysis import (
15 MODEL_FAILURE_RETRY_DELAY,
16 SonicAnalysisProvider,
17 SonicSessionData,
18)
19from music_assistant.providers.sonic_analysis.clap_prompts import (
20 CALIBRATION,
21 SCALAR_PROMPT_PAIRS,
22)
23
24
25def _make_provider() -> tuple[SonicAnalysisProvider, MagicMock]:
26 """Stub provider with the prompt order needed by _run_live_clap_if_eligible."""
27 p = SonicAnalysisProvider.__new__(SonicAnalysisProvider)
28 fake_logger = MagicMock()
29 p.logger = fake_logger
30 p._clap_prompt_order = list(SCALAR_PROMPT_PAIRS.items())
31 return p, fake_logger
32
33
34def _make_session(target_starts: list[int] | None = None) -> SonicSessionData:
35 starts = list(target_starts) if target_starts is not None else []
36 return SonicSessionData(
37 streamdetails=MagicMock(),
38 audio_format=MagicMock(),
39 clap_target_starts=starts,
40 clap_target_buffers=[[] for _ in starts],
41 clap_target_complete=[False] * len(starts),
42 )
43
44
45@pytest.mark.asyncio
46async def test_no_targets_short_circuits_silently() -> None:
47 """A session with no planned targets returns immediately, no log noise."""
48 p, fake_logger = _make_provider()
49 session = _make_session(target_starts=[])
50 analysis = AudioAnalysisData()
51
52 await p._run_live_clap_if_eligible(session, analysis)
53
54 assert analysis.danceability is None
55 assert analysis.valence is None
56 assert analysis.arousal is None
57 assert analysis.instrumentalness is None
58 assert analysis.acousticness is None
59 assert analysis.speechiness is None
60 assert analysis.extra_data is None or "clap_embedding" not in (analysis.extra_data or {})
61 fake_logger.warning.assert_not_called()
62
63
64@pytest.mark.asyncio
65async def test_no_completions_raises_retryable() -> None:
66 """Targets planned but zero windows completed â retryable failure, no scalar updates."""
67 p, fake_logger = _make_provider()
68 session = _make_session(target_starts=[0, 100, 200])
69 # No tasks added, no completed_count incremented
70 analysis = AudioAnalysisData()
71
72 before = utc()
73 with pytest.raises(AudioAnalysisError) as excinfo:
74 await p._run_live_clap_if_eligible(session, analysis)
75
76 assert excinfo.value.retry_at is not None
77 assert excinfo.value.retry_at >= before + MODEL_FAILURE_RETRY_DELAY
78
79 assert analysis.danceability is None
80 assert analysis.valence is None
81 assert analysis.arousal is None
82 assert analysis.instrumentalness is None
83 assert analysis.acousticness is None
84 assert analysis.speechiness is None
85 assert analysis.extra_data is None or "clap_embedding" not in (analysis.extra_data or {})
86 fake_logger.warning.assert_called_once()
87
88
89@pytest.mark.asyncio
90async def test_partial_completions_raises_retryable() -> None:
91 """Some but not all planned windows completed â retryable failure, nothing written."""
92 p, fake_logger = _make_provider()
93 session = _make_session(target_starts=[0, 100, 200])
94
95 n_pairs = len(SCALAR_PROMPT_PAIRS)
96 session.clap_completed_count = 2
97 session.clap_sum_embedding = np.ones(1024, dtype=np.float32)
98 session.clap_sum_similarities = np.zeros(2 * n_pairs, dtype=np.float32)
99
100 analysis = AudioAnalysisData()
101
102 with pytest.raises(AudioAnalysisError, match="2 of 3"):
103 await p._run_live_clap_if_eligible(session, analysis)
104
105 assert analysis.danceability is None
106 assert analysis.extra_data is None or "clap_embedding" not in (analysis.extra_data or {})
107 fake_logger.warning.assert_called_once()
108
109
110@pytest.mark.asyncio
111async def test_mean_pools_and_calibrates_scalars() -> None:
112 """Three completed windows with known sums produce calibrated scalars and L2-normalized embedding."""
113 p, _ = _make_provider()
114 session = _make_session(target_starts=[0, 100, 200])
115
116 n_pairs = len(SCALAR_PROMPT_PAIRS)
117 # Sums equivalent to mean_emb = 0.5 (pre-norm), mean_sim = [1.0, 0.0, 1.0, 0.0, ...]
118 session.clap_completed_count = 3
119 session.clap_sum_embedding = np.full(1024, 1.5, dtype=np.float32) # mean = 0.5
120 raw_sims = np.zeros(2 * n_pairs, dtype=np.float32)
121 for i in range(n_pairs):
122 raw_sims[i * 2] = 3.0 # pos_logit mean = 1.0
123 raw_sims[i * 2 + 1] = 0.0 # neg_logit mean = 0.0
124 session.clap_sum_similarities = raw_sims
125
126 analysis = AudioAnalysisData()
127 await p._run_live_clap_if_eligible(session, analysis)
128
129 # Embedding: pre-norm = 0.5 across 1024 dims; ||v|| = sqrt(1024 * 0.25) = 16.0; normalized = 1/32
130 assert analysis.extra_data is not None
131 emb = np.asarray(analysis.extra_data["clap_embedding"], dtype=np.float32)
132 expected_norm = math.sqrt(1024 * 0.25)
133 expected_value = 0.5 / expected_norm
134 np.testing.assert_array_almost_equal(emb, np.full(1024, expected_value), decimal=5)
135
136 for scalar_name in SCALAR_PROMPT_PAIRS:
137 a, b = CALIBRATION[scalar_name]
138 expected = 1.0 / (1.0 + math.exp(-(a * 1.0 + b)))
139 assert getattr(analysis, scalar_name) == pytest.approx(expected)
140
141
142@pytest.mark.asyncio
143async def test_awaits_pending_tasks() -> None:
144 """Tasks still in flight are awaited before the mean-pool runs."""
145 p, _ = _make_provider()
146 session = _make_session(target_starts=[0])
147
148 completion_event = asyncio.Event()
149
150 async def slow_inference() -> None:
151 await completion_event.wait()
152 n_pairs = len(SCALAR_PROMPT_PAIRS)
153 session.clap_sum_embedding = np.ones(1024, dtype=np.float32)
154 session.clap_sum_similarities = np.zeros(2 * n_pairs, dtype=np.float32)
155 session.clap_completed_count = 1
156
157 task = asyncio.create_task(slow_inference())
158 session.clap_inference_tasks.append(task)
159
160 finalize_task = asyncio.create_task(p._run_live_clap_if_eligible(session, AudioAnalysisData()))
161 # Yield once: finalize should be blocked on the gather
162 await asyncio.sleep(0)
163 assert not finalize_task.done()
164
165 # Release the inference task
166 completion_event.set()
167 await finalize_task
168 assert task.done()
169
170
171@pytest.mark.asyncio
172async def test_cancel_aborts_pending_tasks_and_clears_buffers() -> None:
173 """Cancel cancels in-flight inferences and resets per-window buffers."""
174 p = SonicAnalysisProvider.__new__(SonicAnalysisProvider)
175 p.logger = MagicMock()
176 p._sessions = {}
177
178 session = _make_session(target_starts=[0, 100])
179 session.clap_target_buffers = [
180 [np.zeros(1024, dtype=np.float32)],
181 [np.zeros(2048, dtype=np.float32)],
182 ]
183
184 started = asyncio.Event()
185 cancelled = asyncio.Event()
186
187 async def long_running() -> None:
188 started.set()
189 try:
190 await asyncio.sleep(60)
191 except asyncio.CancelledError:
192 cancelled.set()
193 raise
194
195 task = asyncio.create_task(long_running())
196 session.clap_inference_tasks.append(task)
197 await started.wait()
198
199 p._sessions["sess"] = session
200 await p.cancel("sess")
201
202 # Give the cancellation a tick to propagate
203 await asyncio.sleep(0)
204 assert task.cancelled() or cancelled.is_set()
205 assert session.clap_target_buffers == []
206