music-assistant-server

1"""Tests for the CLAP model load that gates SonicAnalysisProvider setup."""
2
3from __future__ import annotations
4
5import asyncio
6from collections.abc import Generator
7from typing import Any
8from unittest.mock import AsyncMock, MagicMock, patch
9
10import httpx
11import pytest
12from music_assistant_models.enums import ContentType
13from music_assistant_models.errors import SetupFailedError, UnsupportedSystemError
14from music_assistant_models.media_items import AudioFormat
15
16from music_assistant.constants import CONF_LOG_LEVEL
17from music_assistant.providers.sonic_analysis import (
18    CLAP_SAMPLING_FAST,
19    SonicAnalysisProvider,
20)
21
22SETUP_TASK_ID = "sonic_analysis.model_setup.instance-1"
23
24# --- Helpers ---
25
26
27class _FakeMass:
28    """
29    Stand-in for MusicAssistant reproducing the ``create_task`` dedupe setup relies on.
30
31    A MagicMock would return one from ``create_task``, leaving the coroutine unawaited.
32    """
33
34    def __init__(self) -> None:
35        self.streams = MagicMock()
36        self.cache = MagicMock()
37        self.tracked: dict[str, asyncio.Task[Any]] = {}
38        self.task_ids: list[str | None] = []
39
40    def create_task(
41        self, target: Any, *args: Any, task_id: str | None = None, **kwargs: Any
42    ) -> asyncio.Task[Any]:
43        """Return the live task registered under task_id, else start and track a new one."""
44        self.task_ids.append(task_id)
45        if task_id and (existing := self.tracked.get(task_id)) and not existing.done():
46            target.close()
47            return existing
48        task: asyncio.Task[Any] = asyncio.ensure_future(target)
49        # Mirrors create_task's done callback, which retrieves the exception.
50        task.add_done_callback(lambda t: t.cancelled() or t.exception())
51        if task_id:
52            self.tracked[task_id] = task
53        return task
54
55    async def drain(self) -> None:
56        """Cancel and await every tracked task, so no test leaks one."""
57        for task in self.tracked.values():
58            task.cancel()
59        await asyncio.gather(*self.tracked.values(), return_exceptions=True)
60
61
62def _make_provider(mass: _FakeMass | None = None) -> SonicAnalysisProvider:
63    """
64    Construct a SonicAnalysisProvider with mocked MA infrastructure.
65
66    :param mass: Shared stand-in, for tests needing two providers on one task registry.
67    """
68    manifest = MagicMock()
69    manifest.domain = "sonic_analysis"
70
71    config = MagicMock()
72    config.instance_id = "instance-1"
73    config.get_value = MagicMock(
74        side_effect=lambda key, *_a, **_kw: (
75            "GLOBAL" if key == CONF_LOG_LEVEL else CLAP_SAMPLING_FAST
76        )
77    )
78
79    provider = SonicAnalysisProvider(mass or _FakeMass(), manifest, config)  # type: ignore[arg-type]
80    provider.logger = MagicMock()
81    return provider
82
83
84def _fake_models() -> tuple[Any, Any, list[tuple[str, tuple[str, str]]]]:
85    """Return a stand-in for what a completed CLAP load hands back."""
86    return MagicMock(name="clap_model"), MagicMock(name="text_embeddings"), []
87
88
89def _make_audio_format() -> AudioFormat:
90    """Return a real AudioFormat for 16-bit mono PCM."""
91    return AudioFormat(
92        content_type=ContentType.PCM_S16LE, sample_rate=22050, bit_depth=16, channels=1
93    )
94
95
96def _make_streamdetails(item_id: str = "track-1", duration: float | None = 60.0) -> MagicMock:
97    """Return a minimal streamdetails mock."""
98    sd = MagicMock()
99    sd.item_id = item_id
100    sd.provider = "test_provider"
101    sd.duration = duration
102    return sd
103
104
105@pytest.fixture(autouse=True)
106def _stub_ml_inference_gate() -> Generator[None]:
107    """Stub the hardware gate so these unit tests never spawn the real capability probe."""
108    with patch(
109        "music_assistant.providers.sonic_analysis.verify_system_meets_requirements",
110        new=AsyncMock(),
111    ):
112        yield
113
114
115# --- handle_async_init ---
116
117
118@pytest.mark.asyncio
119async def test_completed_load_populates_state() -> None:
120    """A load that finishes in time must populate model state and mark models loaded."""
121    provider = _make_provider()
122    model, embeddings, prompt_order = _fake_models()
123
124    with patch.object(provider, "_load_clap", return_value=(model, embeddings, prompt_order)):
125        await provider.handle_async_init()
126
127    assert provider._clap_model is model
128    assert provider._clap_text_embeddings is embeddings
129    assert provider._clap_prompt_order == prompt_order
130    assert provider._models_loaded is True
131
132
133@pytest.mark.asyncio
134async def test_load_failure_propagates() -> None:
135    """Load failures must propagate, so the provider is left unavailable."""
136    provider = _make_provider()
137
138    with (
139        patch.object(provider, "_load_clap", side_effect=RuntimeError("checkpoint is corrupt")),
140        pytest.raises(RuntimeError, match="checkpoint is corrupt"),
141    ):
142        await provider.handle_async_init()
143
144    assert provider._clap_model is None
145    assert provider._models_loaded is False
146
147
148@pytest.mark.asyncio
149async def test_load_is_offloaded_to_a_thread() -> None:
150    """The download and model build must run off the event loop."""
151    provider = _make_provider()
152
153    with patch(
154        "music_assistant.providers.sonic_analysis.asyncio.to_thread",
155        new=AsyncMock(return_value=_fake_models()),
156    ) as to_thread_mock:
157        await provider.handle_async_init()
158
159    # ``==`` not ``is``: each attribute access yields a fresh bound-method object.
160    assert to_thread_mock.call_args.args[0] == provider._load_clap
161
162
163@pytest.mark.asyncio
164async def test_unsupported_system_fails_before_any_download() -> None:
165    """An unsupported host must fail before the checkpoint is fetched."""
166    provider = _make_provider()
167
168    with (
169        patch(
170            "music_assistant.providers.sonic_analysis.verify_system_meets_requirements",
171            side_effect=UnsupportedSystemError("unsupported system"),
172        ),
173        patch.object(SonicAnalysisProvider, "_load_clap") as load_clap_mock,
174        pytest.raises(UnsupportedSystemError),
175    ):
176        await provider.handle_async_init()
177
178    load_clap_mock.assert_not_called()
179
180
181@pytest.mark.asyncio
182async def test_slow_load_fails_setup_but_keeps_running() -> None:
183    """A load that outlives the grace period fails setup but keeps running."""
184    mass = _FakeMass()
185    provider = _make_provider(mass)
186
187    with (
188        patch("music_assistant.providers.sonic_analysis.MODEL_SETUP_GRACE_SECONDS", 0.05),
189        patch(
190            "music_assistant.providers.sonic_analysis.asyncio.to_thread",
191            new=lambda *_a: asyncio.Event().wait(),
192        ),
193        pytest.raises(SetupFailedError) as exc_info,
194    ):
195        await provider.handle_async_init()
196
197    assert exc_info.value.translation_key == "model_setup_pending"
198    assert exc_info.value.translation_owner == "provider.sonic_analysis"
199    assert provider._models_loaded is False
200    assert not mass.tracked[SETUP_TASK_ID].done(), "the load must survive the timeout"
201
202    await mass.drain()
203
204
205@pytest.mark.asyncio
206async def test_retry_joins_the_running_load_and_gets_its_result() -> None:
207    """A retry joins the running load and comes away with its result."""
208    mass = _FakeMass()
209    release = asyncio.Event()
210    models = _fake_models()
211    load_calls = 0
212
213    async def _blocked_load(*_args: Any) -> tuple[Any, Any, list[Any]]:
214        nonlocal load_calls
215        load_calls += 1
216        await release.wait()
217        return models
218
219    with patch("music_assistant.providers.sonic_analysis.asyncio.to_thread", new=_blocked_load):
220        first = _make_provider(mass)
221        with (
222            patch("music_assistant.providers.sonic_analysis.MODEL_SETUP_GRACE_SECONDS", 0.05),
223            pytest.raises(SetupFailedError),
224        ):
225            await first.handle_async_init()
226
227        second = _make_provider(mass)
228        # the load lands partway through the retry's own grace period
229        asyncio.get_running_loop().call_soon(release.set)
230        await second.handle_async_init()
231
232    assert load_calls == 1, "the retry must not start a second load"
233    assert mass.task_ids == [SETUP_TASK_ID, SETUP_TASK_ID], "the key must be stable"
234    assert second._clap_model is models[0]
235    assert second._models_loaded is True
236
237
238# --- _load_clap ---
239
240
241@pytest.mark.parametrize(
242    "err",
243    [OSError("disk full"), httpx.ConnectError("connection refused")],
244    ids=["oserror", "httpx_connect"],
245)
246def test_download_failures_become_retryable_setup_errors(err: Exception) -> None:
247    """A failed download must reach MA as a typed error, which is what gets it retried."""
248    provider = _make_provider()
249
250    with (
251        patch.object(provider, "_try_load_cached_prompt_embeddings", return_value=MagicMock()),
252        patch("music_assistant.providers.sonic_analysis.vendored_clap.CLAP", side_effect=err),
253        pytest.raises(SetupFailedError) as exc_info,
254    ):
255        provider._load_clap()
256
257    assert exc_info.value.translation_key == "model_assets_download_failed"
258    assert exc_info.value.__cause__ is err
259
260
261def test_missing_prompt_embeddings_fail_instead_of_downloading_a_text_encoder() -> None:
262    """Missing or stale prompt embeddings must fail, not fall back to the text encoder."""
263    provider = _make_provider()
264
265    with (
266        patch.object(provider, "_try_load_cached_prompt_embeddings", return_value=None),
267        patch("music_assistant.providers.sonic_analysis.vendored_clap.CLAP") as clap_cls,
268        pytest.raises(UnsupportedSystemError) as exc_info,
269    ):
270        provider._load_clap()
271
272    # the exact type matters: it is the one setup error MA does not retry, and no
273    # retry can recreate a shipped file
274    assert type(exc_info.value) is UnsupportedSystemError
275    assert exc_info.value.translation_key == "prompt_embeddings_unavailable"
276    clap_cls.assert_not_called()
277
278
279# --- _start_analysis gating ---
280
281
282@pytest.mark.asyncio
283async def test_start_analysis_declines_while_clap_is_unavailable() -> None:
284    """``_start_analysis`` must decline tracks while CLAP is unavailable."""
285    provider = _make_provider()
286    provider._clap_model = None
287
288    result = await provider._start_analysis(
289        "session-skip", _make_streamdetails("skip-me"), _make_audio_format()
290    )
291
292    assert result is False
293    assert provider._sessions == {}
294
295
296@pytest.mark.asyncio
297async def test_start_analysis_proceeds_when_clap_loaded() -> None:
298    """``_start_analysis`` must create a session when CLAP is available."""
299    provider = _make_provider()
300    provider._clap_model = MagicMock(name="clap_model")
301
302    result = await provider._start_analysis(
303        "session-ok", _make_streamdetails("go-ahead"), _make_audio_format()
304    )
305
306    assert result is True
307    assert "session-ok" in provider._sessions
308
309
310@pytest.mark.asyncio
311@pytest.mark.parametrize("duration", [None, 0, 0.0])
312async def test_start_analysis_declines_without_duration(duration: float | None) -> None:
313    """``_start_analysis`` must decline tracks without a usable duration."""
314    provider = _make_provider()
315    provider._clap_model = MagicMock(name="clap_model")
316
317    result = await provider._start_analysis(
318        "session-no-duration", _make_streamdetails("no-duration", duration), _make_audio_format()
319    )
320
321    assert result is False
322    assert provider._sessions == {}
323