/
/
/
1"""Unit tests for Yandex Music streaming quality selection."""
2
3from __future__ import annotations
4
5import unittest.mock
6from typing import TYPE_CHECKING, Any
7
8import pytest
9from aiohttp import ClientPayloadError
10from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
11from music_assistant_models.enums import ContentType, StreamType
12from music_assistant_models.errors import MediaNotFoundError
13from music_assistant_models.media_items import AudioFormat
14from music_assistant_models.streamdetails import StreamDetails
15
16from music_assistant.providers.yandex_music.constants import (
17 QUALITY_BALANCED,
18 QUALITY_EFFICIENT,
19 QUALITY_HIGH,
20 QUALITY_SUPERB,
21)
22from music_assistant.providers.yandex_music.streaming import YandexMusicStreamingManager
23
24if TYPE_CHECKING:
25 from tests.providers.yandex_music.conftest import (
26 StreamingProviderStub,
27 StreamingProviderStubWithTracking,
28 )
29
30
31def _make_download_info(
32 codec: str,
33 bitrate_in_kbps: int,
34 direct_link: str = "https://example.com/track",
35) -> Any:
36 """Build DownloadInfo-like object."""
37 return type(
38 "DownloadInfo",
39 (),
40 {
41 "codec": codec,
42 "bitrate_in_kbps": bitrate_in_kbps,
43 "direct_link": direct_link,
44 },
45 )()
46
47
48@pytest.fixture
49def streaming_manager(
50 streaming_provider_stub: StreamingProviderStub,
51) -> YandexMusicStreamingManager:
52 """Create streaming manager with real stub (no Mock)."""
53 return YandexMusicStreamingManager(streaming_provider_stub) # type: ignore[arg-type]
54
55
56@pytest.fixture
57def streaming_manager_with_tracking(
58 streaming_provider_stub_with_tracking: StreamingProviderStubWithTracking,
59) -> YandexMusicStreamingManager:
60 """Create streaming manager with tracking logger for assertions."""
61 return YandexMusicStreamingManager(streaming_provider_stub_with_tracking) # type: ignore[arg-type]
62
63
64def test_select_best_quality_lossless_returns_flac(
65 streaming_manager: YandexMusicStreamingManager,
66) -> None:
67 """When preferred_quality is 'lossless' and list has MP3 and FLAC, FLAC is selected."""
68 mp3 = _make_download_info("mp3", 320, "https://example.com/track.mp3")
69 flac = _make_download_info("flac", 0, "https://example.com/track.flac")
70 download_infos = [mp3, flac]
71
72 result = streaming_manager._select_best_quality(download_infos, QUALITY_SUPERB)
73
74 assert result is not None
75 assert result.codec == "flac"
76 assert result.direct_link == "https://example.com/track.flac"
77
78
79def test_select_best_quality_balanced_falls_back_to_highest(
80 streaming_manager: YandexMusicStreamingManager,
81) -> None:
82 """When preferred is 'balanced' and no option in 128-256kbps range, highest bitrate is used."""
83 mp3 = _make_download_info("mp3", 320, "https://example.com/track.mp3")
84 flac = _make_download_info("flac", 0, "https://example.com/track.flac")
85 download_infos = [mp3, flac]
86
87 result = streaming_manager._select_best_quality(download_infos, QUALITY_BALANCED)
88
89 assert result is not None
90 assert result.codec == "mp3"
91 assert result.bitrate_in_kbps == 320
92
93
94def test_select_best_quality_label_lossless_flac_returns_flac(
95 streaming_manager: YandexMusicStreamingManager,
96) -> None:
97 """When preferred_quality is UI label 'Lossless (FLAC)', FLAC is selected."""
98 mp3 = _make_download_info("mp3", 320, "https://example.com/track.mp3")
99 flac = _make_download_info("flac", 0, "https://example.com/track.flac")
100 download_infos = [mp3, flac]
101
102 result = streaming_manager._select_best_quality(download_infos, "Lossless (FLAC)")
103
104 assert result is not None
105 assert result.codec == "flac"
106
107
108def test_select_best_quality_lossless_no_flac_returns_fallback(
109 streaming_manager_with_tracking: YandexMusicStreamingManager,
110) -> None:
111 """When lossless requested but no FLAC in list, returns best available (fallback)."""
112 mp3 = _make_download_info("mp3", 320, "https://example.com/track.mp3")
113 download_infos = [mp3]
114
115 result = streaming_manager_with_tracking._select_best_quality(download_infos, QUALITY_SUPERB)
116
117 assert result is not None
118 assert result.codec == "mp3"
119 assert streaming_manager_with_tracking.provider.logger._warning_count == 1 # type: ignore[attr-defined]
120
121
122def test_select_best_quality_empty_list_returns_none(
123 streaming_manager: YandexMusicStreamingManager,
124) -> None:
125 """Empty download_infos returns None."""
126 result = streaming_manager._select_best_quality([], QUALITY_SUPERB)
127 assert result is None
128
129
130def test_select_best_quality_none_preferred_returns_highest_bitrate(
131 streaming_manager: YandexMusicStreamingManager,
132) -> None:
133 """When preferred_quality is None, returns highest bitrate."""
134 mp3 = _make_download_info("mp3", 320, "https://example.com/track.mp3")
135 flac = _make_download_info("flac", 0, "https://example.com/track.flac")
136 download_infos = [mp3, flac]
137
138 result = streaming_manager._select_best_quality(download_infos, None)
139
140 assert result is not None
141 assert result.codec == "mp3"
142 assert result.bitrate_in_kbps == 320
143
144
145def test_get_content_type_flac_mp4_returns_mp4_container_with_flac_codec(
146 streaming_manager: YandexMusicStreamingManager,
147) -> None:
148 """flac-mp4 codec from get-file-info is mapped to MP4 container with FLAC codec."""
149 assert streaming_manager._get_content_type("flac-mp4") == (ContentType.MP4, ContentType.FLAC)
150 assert streaming_manager._get_content_type("FLAC-MP4") == (ContentType.MP4, ContentType.FLAC)
151
152
153def test_get_content_type_flac_returns_flac_container_with_unknown_codec(
154 streaming_manager: YandexMusicStreamingManager,
155) -> None:
156 """Plain FLAC codec is mapped to FLAC container with UNKNOWN codec."""
157 assert streaming_manager._get_content_type("flac") == (ContentType.FLAC, ContentType.UNKNOWN)
158 assert streaming_manager._get_content_type("FLAC") == (ContentType.FLAC, ContentType.UNKNOWN)
159
160
161def test_get_content_type_aac_variants_return_aac(
162 streaming_manager: YandexMusicStreamingManager,
163) -> None:
164 """All AAC codec variants are mapped correctly (MP4 container or plain AAC)."""
165 # Plain AAC variants
166 assert streaming_manager._get_content_type("aac") == (ContentType.AAC, ContentType.UNKNOWN)
167 assert streaming_manager._get_content_type("AAC") == (ContentType.AAC, ContentType.UNKNOWN)
168 assert streaming_manager._get_content_type("he-aac") == (ContentType.AAC, ContentType.UNKNOWN)
169 assert streaming_manager._get_content_type("HE-AAC") == (ContentType.AAC, ContentType.UNKNOWN)
170 # MP4 container variants
171 assert streaming_manager._get_content_type("aac-mp4") == (ContentType.MP4, ContentType.AAC)
172 assert streaming_manager._get_content_type("AAC-MP4") == (ContentType.MP4, ContentType.AAC)
173 assert streaming_manager._get_content_type("he-aac-mp4") == (ContentType.MP4, ContentType.AAC)
174 assert streaming_manager._get_content_type("HE-AAC-MP4") == (ContentType.MP4, ContentType.AAC)
175
176
177# --- Efficient quality tests ---
178
179
180def test_select_best_quality_efficient_prefers_lowest_aac(
181 streaming_manager: YandexMusicStreamingManager,
182) -> None:
183 """Efficient quality prefers lowest bitrate AAC over higher bitrate options."""
184 mp3_320 = _make_download_info("mp3", 320)
185 aac_64 = _make_download_info("aac", 64)
186 aac_192 = _make_download_info("aac", 192)
187
188 result = streaming_manager._select_best_quality([mp3_320, aac_64, aac_192], QUALITY_EFFICIENT)
189
190 assert result is not None
191 assert result.codec == "aac"
192 assert result.bitrate_in_kbps == 64
193
194
195def test_select_best_quality_efficient_aac_mp4_variant(
196 streaming_manager: YandexMusicStreamingManager,
197) -> None:
198 """Efficient quality recognizes aac-mp4 container variant."""
199 mp3_320 = _make_download_info("mp3", 320)
200 aac_mp4_64 = _make_download_info("aac-mp4", 64)
201
202 result = streaming_manager._select_best_quality([mp3_320, aac_mp4_64], QUALITY_EFFICIENT)
203
204 assert result is not None
205 assert result.codec == "aac-mp4"
206 assert result.bitrate_in_kbps == 64
207
208
209def test_select_best_quality_efficient_fallback_to_mp3(
210 streaming_manager: YandexMusicStreamingManager,
211) -> None:
212 """Efficient quality falls back to MP3 when no AAC available."""
213 mp3_128 = _make_download_info("mp3", 128)
214 flac = _make_download_info("flac", 0)
215
216 result = streaming_manager._select_best_quality([mp3_128, flac], QUALITY_EFFICIENT)
217
218 assert result is not None
219 assert result.codec == "mp3"
220
221
222def test_select_best_quality_efficient_fallback_to_lowest(
223 streaming_manager: YandexMusicStreamingManager,
224) -> None:
225 """Efficient quality falls back to lowest bitrate when no AAC/MP3."""
226 flac = _make_download_info("flac", 1411)
227
228 result = streaming_manager._select_best_quality([flac], QUALITY_EFFICIENT)
229
230 assert result is not None
231 assert result.codec == "flac"
232
233
234# --- High quality tests ---
235
236
237def test_select_best_quality_high_prefers_mp3_320(
238 streaming_manager: YandexMusicStreamingManager,
239) -> None:
240 """High quality prefers MP3 with bitrate >= 256kbps."""
241 mp3_320 = _make_download_info("mp3", 320)
242 mp3_128 = _make_download_info("mp3", 128)
243 aac_192 = _make_download_info("aac", 192)
244 flac = _make_download_info("flac", 1411)
245
246 result = streaming_manager._select_best_quality([mp3_320, mp3_128, aac_192, flac], QUALITY_HIGH)
247
248 assert result is not None
249 assert result.codec == "mp3"
250 assert result.bitrate_in_kbps == 320
251
252
253def test_select_best_quality_high_fallback_to_any_mp3(
254 streaming_manager: YandexMusicStreamingManager,
255) -> None:
256 """High quality falls back to any MP3 when no high-bitrate MP3 available."""
257 mp3_128 = _make_download_info("mp3", 128)
258 aac_192 = _make_download_info("aac", 192)
259
260 result = streaming_manager._select_best_quality([mp3_128, aac_192], QUALITY_HIGH)
261
262 assert result is not None
263 assert result.codec == "mp3"
264 assert result.bitrate_in_kbps == 128
265
266
267def test_select_best_quality_high_no_mp3_uses_non_flac(
268 streaming_manager: YandexMusicStreamingManager,
269) -> None:
270 """High quality uses highest non-FLAC when no MP3 available."""
271 aac_192 = _make_download_info("aac", 192)
272 flac = _make_download_info("flac", 1411)
273
274 result = streaming_manager._select_best_quality([aac_192, flac], QUALITY_HIGH)
275
276 assert result is not None
277 assert result.codec == "aac"
278 assert result.bitrate_in_kbps == 192
279
280
281def test_select_best_quality_high_only_flac_returns_flac(
282 streaming_manager: YandexMusicStreamingManager,
283) -> None:
284 """High quality returns FLAC as last resort when nothing else available."""
285 flac = _make_download_info("flac", 1411)
286
287 result = streaming_manager._select_best_quality([flac], QUALITY_HIGH)
288
289 assert result is not None
290 assert result.codec == "flac"
291
292
293# --- Audio params tests ---
294
295
296def test_get_audio_params_flac_mp4(
297 streaming_manager: YandexMusicStreamingManager,
298) -> None:
299 """flac-mp4 returns 48kHz/24bit."""
300 assert streaming_manager._get_audio_params("flac-mp4") == (48000, 24)
301
302
303def test_get_audio_params_flac_mp4_case_insensitive(
304 streaming_manager: YandexMusicStreamingManager,
305) -> None:
306 """flac-mp4 matching is case-insensitive."""
307 assert streaming_manager._get_audio_params("FLAC-MP4") == (48000, 24)
308
309
310def test_get_audio_params_flac(
311 streaming_manager: YandexMusicStreamingManager,
312) -> None:
313 """Plain FLAC returns CD-quality defaults."""
314 assert streaming_manager._get_audio_params("flac") == (44100, 16)
315
316
317def test_get_audio_params_mp3(
318 streaming_manager: YandexMusicStreamingManager,
319) -> None:
320 """MP3 returns CD-quality defaults."""
321 assert streaming_manager._get_audio_params("mp3") == (44100, 16)
322
323
324def test_get_audio_params_none(
325 streaming_manager: YandexMusicStreamingManager,
326) -> None:
327 """None codec returns CD-quality defaults."""
328 assert streaming_manager._get_audio_params(None) == (44100, 16)
329
330
331# --- get_audio_stream tests ---
332
333
334def _make_encrypted_stream_details(
335 key_hex: str,
336 url: str = "https://example.com/encrypted.flac",
337) -> StreamDetails:
338 """Build StreamDetails for encrypted FLAC stream tests."""
339 return StreamDetails(
340 item_id="test_track_123",
341 provider="yandex_music_instance",
342 audio_format=AudioFormat(content_type=ContentType.MP4),
343 stream_type=StreamType.CUSTOM,
344 data={
345 "encrypted_url": url,
346 "decryption_key": key_hex,
347 "codec": "flac-mp4",
348 },
349 )
350
351
352class _MockContent:
353 """Async iterable content for mock HTTP responses."""
354
355 def __init__(self, chunks: list[bytes], *, drop_payload_error: bool = False) -> None:
356 self._chunks = chunks
357 self._drop = drop_payload_error
358
359 async def iter_chunked(self, size: int) -> Any:
360 for chunk in self._chunks:
361 yield chunk
362 if self._drop:
363 raise ClientPayloadError("connection reset by peer")
364
365
366class _MockResponse:
367 """Fake aiohttp ClientResponse for streaming tests."""
368
369 def __init__(
370 self,
371 chunks: list[bytes],
372 *,
373 error: Exception | None = None,
374 drop_payload_error: bool = False,
375 ) -> None:
376 self.content = _MockContent(chunks, drop_payload_error=drop_payload_error)
377 self._error = error
378
379 def raise_for_status(self) -> None:
380 """Raise stored error if set, simulating a non-2xx HTTP response."""
381 if self._error is not None:
382 raise self._error
383
384 async def __aenter__(self) -> _MockResponse:
385 return self
386
387 async def __aexit__(self, *args: object) -> None:
388 pass
389
390
391class _MockHttpSession:
392 """Fake aiohttp ClientSession for streaming tests."""
393
394 def __init__(self, response: _MockResponse) -> None:
395 self._response = response
396
397 def get(self, url: str, **kwargs: object) -> _MockResponse:
398 return self._response
399
400
401class _MultiCallHttpSession:
402 """Fake aiohttp ClientSession returning successive responses and recording calls."""
403
404 def __init__(self, responses: list[_MockResponse]) -> None:
405 self._responses = responses
406 self.calls: list[dict[str, Any]] = []
407
408 def get(self, url: str, **kwargs: object) -> _MockResponse:
409 self.calls.append({"url": url, "headers": kwargs.get("headers", {})})
410 return self._responses[len(self.calls) - 1]
411
412
413async def test_get_audio_stream_invalid_key_length(
414 streaming_manager: YandexMusicStreamingManager,
415) -> None:
416 """Invalid AES key length raises MediaNotFoundError before any HTTP request."""
417 sd = _make_encrypted_stream_details("deadbeef") # 4 bytes â invalid
418
419 with pytest.raises(MediaNotFoundError, match="Unsupported AES key length"):
420 async for _ in streaming_manager.get_audio_stream(sd):
421 pass
422
423
424async def test_get_audio_stream_http_error_raises_media_not_found(
425 streaming_manager: YandexMusicStreamingManager,
426 streaming_provider_stub: StreamingProviderStub,
427) -> None:
428 """HTTP error from encrypted URL is converted to MediaNotFoundError."""
429 key = b"\x00" * 32
430 sd = _make_encrypted_stream_details(key.hex())
431 streaming_provider_stub.mass.http_session = _MockHttpSession(
432 _MockResponse([], error=RuntimeError("403 Forbidden"))
433 )
434
435 with pytest.raises(MediaNotFoundError, match="Failed to fetch encrypted stream"):
436 async for _ in streaming_manager.get_audio_stream(sd):
437 pass
438
439
440async def test_get_audio_stream_decrypts_aes_ctr_correctly(
441 streaming_manager: YandexMusicStreamingManager,
442 streaming_provider_stub: StreamingProviderStub,
443) -> None:
444 """Encrypted stream is decrypted correctly with AES-256-CTR and zero IV."""
445 key = b"\x42" * 32
446 plaintext = b"Hello, Yandex Music FLAC data!\n" * 50
447
448 # Encrypt with the same algorithm used in get_audio_stream
449 nonce_16 = bytes(16)
450 encryptor = Cipher(algorithms.AES(key), modes.CTR(nonce_16)).encryptor()
451 ciphertext = encryptor.update(plaintext) + encryptor.finalize()
452
453 sd = _make_encrypted_stream_details(key.hex())
454 streaming_provider_stub.mass.http_session = _MockHttpSession(_MockResponse([ciphertext]))
455
456 result = b""
457 async for chunk in streaming_manager.get_audio_stream(sd):
458 result += chunk
459
460 assert result == plaintext
461
462
463async def test_get_audio_stream_reconnects_with_range_header(
464 streaming_manager: YandexMusicStreamingManager,
465 streaming_provider_stub: StreamingProviderStub,
466) -> None:
467 """On ClientPayloadError, reconnects with correct Range header and full plaintext restored."""
468 key = b"\x11" * 32
469 # 96 bytes = 6 AES-CTR blocks; split at byte 48 (block boundary)
470 plaintext = b"AAAAAAAAAAAAAAAA" * 3 + b"BBBBBBBBBBBBBBBB" * 3
471
472 nonce_16 = bytes(16)
473 encryptor = Cipher(algorithms.AES(key), modes.CTR(nonce_16)).encryptor()
474 ciphertext = encryptor.update(plaintext) + encryptor.finalize()
475
476 drop_at = 48 # exactly 3 blocks â clean block boundary
477
478 # First request drops after 48 bytes; second serves the remainder
479 first_resp = _MockResponse([ciphertext[:drop_at]], drop_payload_error=True)
480 second_resp = _MockResponse([ciphertext[drop_at:]])
481 session = _MultiCallHttpSession([first_resp, second_resp])
482 streaming_provider_stub.mass.http_session = session
483
484 result = b""
485 with unittest.mock.patch("asyncio.sleep"):
486 async for chunk in streaming_manager.get_audio_stream(
487 _make_encrypted_stream_details(key.hex())
488 ):
489 result += chunk
490
491 assert result == plaintext
492 assert len(session.calls) == 2
493 assert session.calls[0].get("headers") == {}
494 assert session.calls[1]["headers"] == {"Range": f"bytes={drop_at}-"}
495