/
/
/
1"""
2Smoke tests for the optional 1024-dim CLAP usearch index helper.
3
4Round-trip coverage: deterministic labels, add/contains/get, and
5persistence to the sonic_similarity_clap.usearch filename stem under
6the configured storage_path.
7"""
8
9from __future__ import annotations
10
11import asyncio
12import logging
13import threading
14import time
15from contextlib import suppress
16from pathlib import Path
17from typing import TYPE_CHECKING, Any
18
19import numpy as np
20import pytest
21
22from music_assistant.providers.sonic_similarity.clap_index import (
23 CLAP_EMBEDDING_DIM,
24 ClapIndex,
25 derive_label,
26)
27
28if TYPE_CHECKING:
29 from types import SimpleNamespace
30
31
32@pytest.fixture
33def logger() -> logging.Logger:
34 """Provide a quiet logger for the tests."""
35 lg = logging.getLogger("test_clap_index_relocation")
36 lg.addHandler(logging.NullHandler())
37 return lg
38
39
40def _make_mass(tmp_path: Path) -> SimpleNamespace:
41 """Build a minimal mass stand-in exposing storage_path."""
42 from types import SimpleNamespace # noqa: PLC0415
43
44 return SimpleNamespace(storage_path=str(tmp_path))
45
46
47def _unit_vec(seed: int) -> np.ndarray:
48 """Return a deterministic L2-normalized 1024-dim vector."""
49 rng = np.random.default_rng(seed)
50 v = rng.standard_normal(CLAP_EMBEDDING_DIM).astype(np.float32)
51 return v / np.linalg.norm(v)
52
53
54def test_derive_label_is_deterministic_and_provider_aware() -> None:
55 """Same inputs â same label; different providers give different labels."""
56 a = derive_label("spotify", "track123")
57 b = derive_label("spotify", "track123")
58 c = derive_label("tidal", "track123")
59 assert a == b
60 assert a != c
61
62
63@pytest.mark.asyncio
64async def test_round_trip_persists_under_sonic_similarity_stem(
65 tmp_path: Path, logger: logging.Logger
66) -> None:
67 """add/contains/get_embedding_by_item_id round-trip + file persistence check."""
68 idx = ClapIndex(_make_mass(tmp_path), logger) # type: ignore[arg-type]
69 await idx.load()
70
71 vec = _unit_vec(1)
72 await idx.add("spotify", "track_xyz", vec)
73
74 assert idx.contains("spotify", "track_xyz")
75 result = idx.get_embedding_by_item_id("track_xyz")
76 assert result is not None
77 prov, stored_vec = result
78 assert prov == "spotify"
79 assert stored_vec.shape == (CLAP_EMBEDDING_DIM,)
80
81 await idx.save()
82 assert (tmp_path / "sonic_similarity_clap.usearch").exists()
83 assert (tmp_path / "sonic_similarity_clap_keys.json").exists()
84
85
86@pytest.mark.asyncio
87async def test_saved_index_reloads_with_its_embeddings(
88 tmp_path: Path, logger: logging.Logger
89) -> None:
90 """A saved index comes back intact through a fresh instance."""
91 idx = ClapIndex(_make_mass(tmp_path), logger) # type: ignore[arg-type]
92 await idx.load()
93 vectors = {f"track{n}": _unit_vec(n) for n in range(20)}
94 for item_id, vec in vectors.items():
95 await idx.add("spotify", item_id, vec)
96 await idx.save()
97
98 reloaded = ClapIndex(_make_mass(tmp_path), logger) # type: ignore[arg-type]
99 await reloaded.load()
100
101 assert len(reloaded) == len(vectors)
102 for item_id, vec in vectors.items():
103 stored = reloaded.get_embedding("spotify", item_id)
104 assert stored is not None
105 np.testing.assert_allclose(stored, vec, atol=1e-3)
106
107
108@pytest.mark.asyncio
109async def test_get_embedding_by_item_id_missing_returns_none(
110 tmp_path: Path, logger: logging.Logger
111) -> None:
112 """Lookup of an unknown item returns None, not an exception."""
113 idx = ClapIndex(_make_mass(tmp_path), logger) # type: ignore[arg-type]
114 await idx.load()
115 assert idx.get_embedding_by_item_id("not_in_index") is None
116
117
118@pytest.mark.asyncio
119async def test_get_embedding_o1_lookup_matches_provider(
120 tmp_path: Path, logger: logging.Logger
121) -> None:
122 """get_embedding returns the vector for the right provider and None otherwise."""
123 idx = ClapIndex(_make_mass(tmp_path), logger) # type: ignore[arg-type]
124 await idx.load()
125 vec = _unit_vec(7)
126 await idx.add("spotify", "track_xyz", vec)
127
128 stored = idx.get_embedding("spotify", "track_xyz")
129 assert stored is not None
130 assert stored.shape == (CLAP_EMBEDDING_DIM,)
131 np.testing.assert_allclose(stored, vec, atol=1e-3)
132
133 # Same item_id under a different provider derives a different label â miss.
134 assert idx.get_embedding("tidal", "track_xyz") is None
135 assert idx.get_embedding("spotify", "not_in_index") is None
136
137
138@pytest.mark.asyncio
139async def test_save_does_not_leave_tmp_file_behind(tmp_path: Path, logger: logging.Logger) -> None:
140 """The atomic-rename path consumes the .tmp file; nothing lingers after success."""
141 idx = ClapIndex(_make_mass(tmp_path), logger) # type: ignore[arg-type]
142 await idx.load()
143 await idx.add("spotify", "track1", _unit_vec(1))
144
145 await idx.save()
146
147 assert not (tmp_path / "sonic_similarity_clap_keys.json.tmp").exists()
148 assert (tmp_path / "sonic_similarity_clap_keys.json").exists()
149
150
151@pytest.mark.asyncio
152async def test_save_writes_keys_before_index_so_a_crash_doesnt_orphan_labels(
153 tmp_path: Path,
154 logger: logging.Logger,
155 monkeypatch: pytest.MonkeyPatch,
156) -> None:
157 """If the binary index save fails mid-flight, the keys file is already up-to-date."""
158 idx = ClapIndex(_make_mass(tmp_path), logger) # type: ignore[arg-type]
159 await idx.load()
160 await idx.add("spotify", "track1", _unit_vec(1))
161
162 # Force the index save to fail after the keys file has already been written.
163 def _boom() -> Any:
164 raise RuntimeError("disk full mid-save")
165
166 monkeypatch.setattr(idx._index, "save", _boom)
167
168 with pytest.raises(RuntimeError, match="disk full"):
169 await idx.save()
170
171 # Keys file was renamed in before the index step blew up; no .tmp lingers.
172 keys_path = tmp_path / "sonic_similarity_clap_keys.json"
173 assert keys_path.exists()
174 assert "track1" in keys_path.read_text(encoding="utf-8")
175 assert not (tmp_path / "sonic_similarity_clap_keys.json.tmp").exists()
176
177
178@pytest.mark.asyncio
179async def test_close_waits_for_a_save_whose_task_was_cancelled(
180 tmp_path: Path,
181 logger: logging.Logger,
182 monkeypatch: pytest.MonkeyPatch,
183) -> None:
184 """Server shutdown cancels the save task; its worker keeps writing and must finish alone."""
185 idx = ClapIndex(_make_mass(tmp_path), logger) # type: ignore[arg-type]
186 await idx.load()
187 await idx.add("spotify", "track1", _unit_vec(1))
188
189 real_save = idx._index.save
190 in_save = threading.Event()
191 release = threading.Event()
192 writers_lock = threading.Lock()
193 writers = 0
194
195 def _slow_save() -> Any:
196 nonlocal writers
197 with writers_lock:
198 writers += 1
199 in_save.set()
200 # parked here rather than sleeping, so the assertion below can never
201 # be decided by how fast the machine happens to be
202 assert release.wait(10), "close() never let the save finish"
203 return real_save()
204
205 monkeypatch.setattr(idx._index, "save", _slow_save)
206
207 save_task = asyncio.create_task(idx.save())
208 assert await asyncio.to_thread(in_save.wait, 10)
209 save_task.cancel()
210 with suppress(asyncio.CancelledError):
211 await save_task
212
213 # the unload that follows task cancellation on shutdown
214 close_task = asyncio.create_task(idx.close())
215
216 # an unguarded close would start writing here; the lock keeps it queued
217 # behind the worker still parked above
218 deadline = time.monotonic() + 0.25
219 while time.monotonic() < deadline and writers < 2:
220 await asyncio.sleep(0.01)
221 assert writers == 1
222
223 release.set()
224 await close_task
225
226 keys_path = tmp_path / "sonic_similarity_clap_keys.json"
227 assert "track1" in keys_path.read_text(encoding="utf-8")
228 assert idx._index is None
229
230
231@pytest.mark.asyncio
232async def test_lookup_survives_inserts_from_the_rebuild_worker(
233 tmp_path: Path, logger: logging.Logger
234) -> None:
235 """A lookup on the event loop must not break while a rebuild inserts from its thread."""
236
237 class _StubIndex:
238 """Stand-in so the insert loop costs the map, not usearch."""
239
240 def __contains__(self, label: int) -> bool:
241 return False
242
243 def add(self, label: int, vec: np.ndarray) -> None:
244 """Accept the vector and discard it."""
245
246 idx = ClapIndex(_make_mass(tmp_path), logger) # type: ignore[arg-type]
247 idx._index = _StubIndex()
248 idx._reverse = {label: ("spotify", f"track{label}") for label in range(40_000)}
249
250 stop = threading.Event()
251 vec = _unit_vec(1)
252
253 def _insert_until_stopped() -> None:
254 for label in range(10_000_000, 10_050_000):
255 if stop.is_set():
256 return
257 idx._add_sync(label, vec, "spotify", f"new{label}")
258
259 worker = threading.Thread(target=_insert_until_stopped, daemon=True)
260 worker.start()
261 try:
262 for _ in range(20):
263 assert idx.get_embedding_by_item_id("absent_track") is None
264 finally:
265 stop.set()
266 worker.join(timeout=5)
267
268
269@pytest.mark.asyncio
270async def test_add_landing_after_release_leaves_no_phantom_entry(
271 tmp_path: Path, logger: logging.Logger
272) -> None:
273 """An insert that reaches the worker after teardown must not record a key without a vector."""
274 idx = ClapIndex(_make_mass(tmp_path), logger) # type: ignore[arg-type]
275 await idx.load()
276 await idx.close()
277
278 # the worker an in-flight add() would have reached once the index was released
279 applied = idx._add_sync(derive_label("spotify", "late"), _unit_vec(1), "spotify", "late")
280
281 assert applied is False
282 assert len(idx) == 0
283 assert not idx.contains("spotify", "late")
284
285
286@pytest.mark.asyncio
287async def test_corrupt_index_with_valid_keys_drops_keys_to_avoid_phantom_contains(
288 tmp_path: Path, logger: logging.Logger
289) -> None:
290 """A corrupt .usearch file alongside a valid keys file must clear both, not strand contains()."""
291 keys_path = tmp_path / "sonic_similarity_clap_keys.json"
292 index_path = tmp_path / "sonic_similarity_clap.usearch"
293 # Plausible-shaped keys file referencing a label that no index backs.
294 stale_label = derive_label("spotify", "phantom_track")
295 keys_path.write_text(f'{{"{stale_label}": ["spotify", "phantom_track"]}}', encoding="utf-8")
296 # Garbage bytes â usearch.load() will raise.
297 index_path.write_bytes(b"not a usearch index")
298
299 idx = ClapIndex(_make_mass(tmp_path), logger) # type: ignore[arg-type]
300 await idx.load()
301
302 assert len(idx) == 0
303 assert not idx.contains("spotify", "phantom_track")
304 assert not keys_path.exists()
305 assert not index_path.exists()
306
307
308@pytest.mark.asyncio
309async def test_missing_index_with_orphan_keys_drops_keys_to_avoid_phantom_contains(
310 tmp_path: Path, logger: logging.Logger
311) -> None:
312 """A keys file present without its index is meaningless â must be cleaned up on load."""
313 keys_path = tmp_path / "sonic_similarity_clap_keys.json"
314 stale_label = derive_label("tidal", "orphan_track")
315 keys_path.write_text(f'{{"{stale_label}": ["tidal", "orphan_track"]}}', encoding="utf-8")
316
317 idx = ClapIndex(_make_mass(tmp_path), logger) # type: ignore[arg-type]
318 await idx.load()
319
320 assert len(idx) == 0
321 assert not idx.contains("tidal", "orphan_track")
322 assert not keys_path.exists()
323