/
/
/
1"""Unit tests for sonic similarity pure functions."""
2
3import pytest
4
5from music_assistant.providers.sonic_similarity.similarity import (
6 apply_mmr,
7 combine_seeds_centroid,
8 expand_recursive,
9 merge_union_results,
10)
11
12
13class TestCombineSeedsCentroid:
14 """Tests for centroid blending of seed signatures."""
15
16 def test_single_seed_returns_itself(self) -> None:
17 """A single seed with no weights returns the same vector."""
18 seed = [1.0, 2.0, 3.0]
19 result = combine_seeds_centroid([seed])
20 assert result == pytest.approx(seed)
21
22 def test_two_seeds_equal_weight(self) -> None:
23 """Two seeds with no weights returns their average."""
24 a = [0.0, 4.0, 6.0]
25 b = [2.0, 0.0, 2.0]
26 result = combine_seeds_centroid([a, b])
27 assert result == pytest.approx([1.0, 2.0, 4.0])
28
29 def test_two_seeds_weighted(self) -> None:
30 """Weighted centroid biases toward the heavier seed."""
31 a = [0.0, 0.0, 0.0]
32 b = [10.0, 10.0, 10.0]
33 result = combine_seeds_centroid([a, b], weights=[1.0, 3.0])
34 assert result == pytest.approx([7.5, 7.5, 7.5])
35
36 def test_weights_normalized(self) -> None:
37 """Weights [2.0, 2.0] produces same result as [0.5, 0.5]."""
38 a = [0.0, 0.0]
39 b = [10.0, 10.0]
40 r1 = combine_seeds_centroid([a, b], weights=[2.0, 2.0])
41 r2 = combine_seeds_centroid([a, b], weights=[0.5, 0.5])
42 assert r1 == pytest.approx(r2)
43
44 def test_empty_seeds_raises(self) -> None:
45 """Empty seed list raises ValueError."""
46 with pytest.raises(ValueError, match="at least one seed"):
47 combine_seeds_centroid([])
48
49 def test_mismatched_weights_raises(self) -> None:
50 """Weights length must match seeds length."""
51 with pytest.raises(ValueError, match="weights length"):
52 combine_seeds_centroid([[1.0, 2.0]], weights=[1.0, 2.0])
53
54
55class TestMergeUnionResults:
56 """Tests for union-mode neighborhood merging."""
57
58 def test_single_neighborhood(self) -> None:
59 """Single neighborhood passes through unchanged."""
60 results = [("track_a", 0.1), ("track_b", 0.5)]
61 merged = merge_union_results([results])
62 assert merged == [("track_a", 0.1), ("track_b", 0.5)]
63
64 def test_two_neighborhoods_dedup_keeps_best(self) -> None:
65 """Duplicate tracks across neighborhoods keep lowest distance."""
66 n1 = [("track_a", 0.3), ("track_b", 0.5)]
67 n2 = [("track_a", 0.1), ("track_c", 0.4)]
68 merged = merge_union_results([n1, n2])
69 merged_dict = dict(merged)
70 assert merged_dict["track_a"] == pytest.approx(0.1)
71 assert "track_b" in merged_dict
72 assert "track_c" in merged_dict
73
74 def test_result_sorted_by_distance(self) -> None:
75 """Merged results are sorted by distance ascending."""
76 n1 = [("a", 0.5)]
77 n2 = [("b", 0.1)]
78 merged = merge_union_results([n1, n2])
79 assert merged[0][0] == "b"
80 assert merged[1][0] == "a"
81
82 def test_empty_neighborhoods(self) -> None:
83 """Empty input returns empty list."""
84 assert merge_union_results([]) == []
85 assert merge_union_results([[]]) == []
86
87
88class TestApplyMMR:
89 """Tests for Maximal Marginal Relevance diversity re-ranking."""
90
91 def test_diversity_zero_preserves_order(self) -> None:
92 """With diversity=0, result order matches pure relevance."""
93 candidates = [
94 ("a", [1.0, 0.0], 0.1),
95 ("b", [0.9, 0.1], 0.2),
96 ("c", [0.5, 0.5], 0.5),
97 ]
98 result = apply_mmr(candidates, [1.0, 0.0], diversity=0.0, limit=3)
99 assert [r[0] for r in result] == ["a", "b", "c"]
100
101 def test_diversity_one_spreads_results(self) -> None:
102 """With diversity=1.0, picks should maximize spread."""
103 candidates = [
104 ("close_1", [1.0, 0.0], 0.1),
105 ("close_2", [0.99, 0.01], 0.11),
106 ("far", [0.0, 1.0], 0.9),
107 ]
108 result = apply_mmr(candidates, [1.0, 0.0], diversity=1.0, limit=3)
109 assert result[0][0] == "close_1"
110 assert result[1][0] == "far"
111
112 def test_limit_respected(self) -> None:
113 """Only limit items are returned."""
114 candidates = [
115 ("a", [1.0], 0.1),
116 ("b", [0.5], 0.2),
117 ("c", [0.0], 0.3),
118 ]
119 result = apply_mmr(candidates, [1.0], diversity=0.0, limit=2)
120 assert len(result) == 2
121
122 def test_empty_candidates(self) -> None:
123 """Empty candidates returns empty."""
124 assert apply_mmr([], [1.0], diversity=0.5, limit=10) == []
125
126
127class TestApplyMMRWeights:
128 """Tests for the preset-weight-aware path of apply_mmr (weights=... shapes the picked set)."""
129
130 @staticmethod
131 def _zero18() -> list[float]:
132 return [0.0] * 18
133
134 @staticmethod
135 def _rhythm_only_weights() -> dict[str, float]:
136 return {
137 "rhythm": 1.0,
138 "loudness": 0.0,
139 "timbre": 0.0,
140 "regularity": 0.0,
141 "mood": 0.0,
142 "tonal": 0.0,
143 "dynamics": 0.0,
144 }
145
146 @staticmethod
147 def _mood_only_weights() -> dict[str, float]:
148 return {
149 "rhythm": 0.0,
150 "loudness": 0.0,
151 "timbre": 0.0,
152 "regularity": 0.0,
153 "mood": 1.0,
154 "tonal": 0.0,
155 "dynamics": 0.0,
156 }
157
158 def test_weights_change_top_pick_diversity_zero(self) -> None:
159 """Switching weight emphasis flips which candidate ranks first."""
160 seed = self._zero18()
161 # FEATURE_GROUPS: rhythm=[0,3), mood=[9,13)
162 # rhythm_match: zero on rhythm, ones elsewhere -> rhythm dist 0, mood dist 1
163 rhythm_match = [0, 0, 0] + [1] * 6 + [1] * 4 + [1] * 5
164 # mood_match: zero on mood, ones elsewhere -> rhythm dist 1, mood dist 0
165 mood_match = [1] * 9 + [0, 0, 0, 0] + [1] * 5
166 candidates = [
167 ("rhythm-cand", [float(v) for v in rhythm_match], 0.5),
168 ("mood-cand", [float(v) for v in mood_match], 0.5),
169 ]
170
171 r1 = apply_mmr(
172 candidates, seed, diversity=0.0, limit=2, weights=self._rhythm_only_weights()
173 )
174 r2 = apply_mmr(candidates, seed, diversity=0.0, limit=2, weights=self._mood_only_weights())
175
176 assert r1[0][0] == "rhythm-cand"
177 assert r2[0][0] == "mood-cand"
178
179 def test_weights_none_matches_default(self) -> None:
180 """weights=None and an explicit None argument produce identical results."""
181 candidates = [
182 ("a", [1.0, 0.0], 0.1),
183 ("b", [0.9, 0.1], 0.2),
184 ("c", [0.5, 0.5], 0.5),
185 ]
186 default = apply_mmr(candidates, [1.0, 0.0], diversity=0.0, limit=3)
187 explicit_none = apply_mmr(candidates, [1.0, 0.0], diversity=0.0, limit=3, weights=None)
188 assert default == explicit_none
189
190 def test_weights_affect_redundancy_when_diversity_positive(self) -> None:
191 """
192 With diversity > 0, weights also affect which 2nd pick MMR diversifies to.
193
194 Setup: rhythm_match_a and rhythm_match_b are both rhythm-close to seed
195 but identical to each other on rhythm (would look 'redundant' under
196 rhythm-only weights). mood_match is rhythm-far but mood-close.
197 Under rhythm-only weights with high diversity, MMR should prefer to
198 avoid picking two rhythm-similar candidates, so the second pick is
199 not the duplicate.
200 """
201 seed = self._zero18()
202 rhythm_a = [0, 0, 0] + [1] * 15 # rhythm dist 0
203 rhythm_b = [0, 0, 0] + [1] * 15 # rhythm dist 0 (duplicate of a)
204 mood_match = [1] * 9 + [0, 0, 0, 0] + [1] * 5 # mood dist 0
205 candidates = [
206 ("rhythm-a", [float(v) for v in rhythm_a], 0.0),
207 ("rhythm-b", [float(v) for v in rhythm_b], 0.0),
208 ("mood-cand", [float(v) for v in mood_match], 1.0),
209 ]
210 result = apply_mmr(
211 candidates,
212 seed,
213 diversity=0.9,
214 limit=2,
215 weights=self._rhythm_only_weights(),
216 )
217 picked = [r[0] for r in result]
218 # First pick: rhythm-a (closest to seed under rhythm-only weights)
219 assert picked[0] in ("rhythm-a", "rhythm-b")
220 # Second pick: not the duplicate -- but under rhythm-only weights,
221 # mood-cand has the same rhythm distance as rhythm-a, so it could be
222 # picked. The duplicate should NOT be picked because it has zero
223 # rhythm-distance from rhythm-a (max redundancy).
224 assert picked[1] != ("rhythm-b" if picked[0] == "rhythm-a" else "rhythm-a")
225
226
227class TestExpandRecursive:
228 """Tests for recursive depth expansion."""
229
230 def test_depth_1_returns_single_generation(self) -> None:
231 """Depth=1 runs the searcher once, all results are generation 0."""
232
233 def searcher(
234 seeds: list[list[float]], # noqa: ARG001
235 seen: set[str], # noqa: ARG001
236 ) -> list[tuple[str, str, list[float], float]]:
237 return [
238 ("a", "prov", [1.0, 0.0], 0.1),
239 ("b", "prov", [0.9, 0.1], 0.2),
240 ]
241
242 results = expand_recursive(
243 initial_seeds=[[1.0, 0.0]], searcher=searcher, depth=1, branch_factor=5
244 )
245 assert len(results) == 2
246 assert all(gen == 0 for _, _, _, _, gen in results)
247
248 def test_depth_2_expands(self) -> None:
249 """Depth=2 uses top branch_factor results from gen 0 as seeds for gen 1."""
250 call_count = 0
251
252 def searcher(
253 seeds: list[list[float]], # noqa: ARG001
254 seen: set[str], # noqa: ARG001
255 ) -> list[tuple[str, str, list[float], float]]:
256 nonlocal call_count
257 call_count += 1
258 if call_count == 1:
259 return [
260 ("a", "p", [1.0, 0.0], 0.1),
261 ("b", "p", [0.8, 0.2], 0.2),
262 ]
263 return [("c", "p", [0.5, 0.5], 0.4)]
264
265 results = expand_recursive(
266 initial_seeds=[[1.0, 0.0]], searcher=searcher, depth=2, branch_factor=2
267 )
268 assert call_count == 2
269 ids = [r[0] for r in results]
270 assert "a" in ids
271 assert "b" in ids
272 assert "c" in ids
273 gen_map = {r[0]: r[4] for r in results}
274 assert gen_map["a"] == 0
275 assert gen_map["b"] == 0
276 assert gen_map["c"] == 1
277
278 def test_deduplication_across_generations(self) -> None:
279 """A track found in gen 0 is not returned again in gen 1."""
280 call_count = 0
281
282 def searcher(
283 seeds: list[list[float]], # noqa: ARG001
284 seen: set[str], # noqa: ARG001
285 ) -> list[tuple[str, str, list[float], float]]:
286 nonlocal call_count
287 call_count += 1
288 if call_count == 1:
289 return [("a", "p", [1.0], 0.1)]
290 return [("b", "p", [0.5], 0.3)]
291
292 results = expand_recursive(
293 initial_seeds=[[1.0]], searcher=searcher, depth=2, branch_factor=1
294 )
295 ids = [r[0] for r in results]
296 assert ids.count("a") == 1
297 assert "b" in ids
298