music-assistant-server

1"""Unit tests for KION Music streaming quality selection."""
2
3from __future__ import annotations
4
5import asyncio
6from typing import TYPE_CHECKING, Any, cast
7
8import pytest
9from aiohttp import ClientPayloadError
10from music_assistant_models.enums import ContentType
11from music_assistant_models.errors import MediaNotFoundError
12from music_assistant_models.media_items import AudioFormat
13from music_assistant_models.streamdetails import StreamDetails
14
15from music_assistant.providers.kion_music.constants import QUALITY_HIGH, QUALITY_LOSSLESS
16from music_assistant.providers.kion_music.streaming import KionMusicStreamingManager
17
18if TYPE_CHECKING:
19    from music_assistant.providers.kion_music.provider import KionMusicProvider
20    from tests.providers.kion_music.conftest import (
21        StreamingProviderStub,
22        StreamingProviderStubWithTracking,
23    )
24
25
26def _make_download_info(
27    codec: str,
28    bitrate_in_kbps: int,
29    direct_link: str = "https://example.com/track",
30) -> Any:
31    """Build DownloadInfo-like object."""
32    return type(
33        "DownloadInfo",
34        (),
35        {
36            "codec": codec,
37            "bitrate_in_kbps": bitrate_in_kbps,
38            "direct_link": direct_link,
39        },
40    )()
41
42
43@pytest.fixture
44def streaming_manager(
45    streaming_provider_stub: StreamingProviderStub,
46) -> KionMusicStreamingManager:
47    """Create streaming manager with real stub (no Mock)."""
48    return KionMusicStreamingManager(cast("KionMusicProvider", streaming_provider_stub))
49
50
51@pytest.fixture
52def streaming_manager_with_tracking(
53    streaming_provider_stub_with_tracking: StreamingProviderStubWithTracking,
54) -> KionMusicStreamingManager:
55    """Create streaming manager with tracking logger for assertions."""
56    return KionMusicStreamingManager(
57        cast("KionMusicProvider", streaming_provider_stub_with_tracking)
58    )
59
60
61def test_select_best_quality_lossless_returns_flac(
62    streaming_manager: KionMusicStreamingManager,
63) -> None:
64    """When preferred_quality is 'lossless' and list has MP3 and FLAC, FLAC is selected."""
65    mp3 = _make_download_info("mp3", 320, "https://example.com/track.mp3")
66    flac = _make_download_info("flac", 0, "https://example.com/track.flac")
67    download_infos = [mp3, flac]
68
69    result = streaming_manager._select_best_quality(download_infos, QUALITY_LOSSLESS)
70
71    assert result is not None
72    assert result.codec == "flac"
73    assert result.direct_link == "https://example.com/track.flac"
74
75
76def test_select_best_quality_high_returns_highest_bitrate(
77    streaming_manager: KionMusicStreamingManager,
78) -> None:
79    """When preferred is 'high' and list has MP3 and FLAC, highest bitrate is selected."""
80    mp3 = _make_download_info("mp3", 320, "https://example.com/track.mp3")
81    flac = _make_download_info("flac", 0, "https://example.com/track.flac")
82    download_infos = [mp3, flac]
83
84    result = streaming_manager._select_best_quality(download_infos, QUALITY_HIGH)
85
86    assert result is not None
87    assert result.codec == "mp3"
88    assert result.bitrate_in_kbps == 320
89
90
91def test_select_best_quality_empty_list_returns_none(
92    streaming_manager: KionMusicStreamingManager,
93) -> None:
94    """Empty download_infos returns None."""
95    result = streaming_manager._select_best_quality([], QUALITY_LOSSLESS)
96    assert result is None
97
98
99def test_select_best_quality_none_preferred_returns_highest_bitrate(
100    streaming_manager: KionMusicStreamingManager,
101) -> None:
102    """When preferred_quality is None, returns highest bitrate."""
103    mp3 = _make_download_info("mp3", 320, "https://example.com/track.mp3")
104    flac = _make_download_info("flac", 0, "https://example.com/track.flac")
105    download_infos = [mp3, flac]
106
107    result = streaming_manager._select_best_quality(download_infos, None)
108
109    assert result is not None
110    assert result.codec == "mp3"
111    assert result.bitrate_in_kbps == 320
112
113
114def _flac_streaminfo_payload(sample_rate: int, bit_depth: int) -> bytes:
115    """
116    Build a 34-byte FLAC STREAMINFO payload for the given sample_rate/bit_depth.
117
118    Layout of bytes 10..14 (4 bytes, big-endian, 32 bits total):
119        sample_rate (20) | channels (3) | bps_minus_1 (5) | total_samples_hi (4)
120    """
121    val = (sample_rate & 0xFFFFF) << 12 | (1 & 0x7) << 9 | ((bit_depth - 1) & 0x1F) << 4
122    return b"\x00" * 10 + val.to_bytes(4, "big") + b"\x00" * 20
123
124
125def _flac_header(sample_rate: int = 44100, bit_depth: int = 16) -> bytes:
126    """Build a valid FLAC file header: magic + block header + STREAMINFO payload."""
127    magic = b"fLaC"
128    block_header = b"\x00\x00\x00\x22"  # type=0 (STREAMINFO), length=34
129    return magic + block_header + _flac_streaminfo_payload(sample_rate, bit_depth)
130
131
132def _mp4_dfla_header(sample_rate: int = 44100, bit_depth: int = 16) -> bytes:
133    """Build an MP4 buffer containing a dfLa box with embedded STREAMINFO."""
134    prefix = b"\x00" * 8  # any 4+ byte prefix so dfl_pos >= 4
135    version_flags = b"\x00" * 4
136    block_header = b"\x00\x00\x00\x22"
137    payload = _flac_streaminfo_payload(sample_rate, bit_depth)
138    return prefix + b"dfLa" + version_flags + block_header + payload
139
140
141def _mp4_mp4a_header(sample_rate: int = 48000, sample_size: int = 16) -> bytes:
142    """Build an MP4 buffer containing an mp4a AudioSampleEntry."""
143    prefix = b"\x00" * 8
144    sr_fixed = (sample_rate << 16) & 0xFFFFFFFF
145    # 0..18: reserved(6) + data_ref(2) + version(2) + revision(2) + vendor(4) + channels(2)
146    # 18..20: sample_size; 20..24: compression_id + packet_size; 24..28: sample_rate (16.16)
147    entry = (
148        b"\x00" * 18 + sample_size.to_bytes(2, "big") + b"\x00" * 4 + sr_fixed.to_bytes(4, "big")
149    )
150    return prefix + b"mp4a" + entry
151
152
153def test_parse_flac_streaminfo_valid() -> None:
154    """Valid FLAC STREAMINFO returns parsed sample_rate and bit_depth."""
155    result = KionMusicStreamingManager._parse_flac_streaminfo(_flac_header(44100, 16))
156    assert result == (44100, 16)
157    result_hires = KionMusicStreamingManager._parse_flac_streaminfo(_flac_header(96000, 24))
158    assert result_hires == (96000, 24)
159
160
161def test_parse_flac_streaminfo_wrong_magic() -> None:
162    """Header without 'fLaC' magic returns (0, 0)."""
163    bad = b"OggS" + _flac_header()[4:]
164    assert KionMusicStreamingManager._parse_flac_streaminfo(bad) == (0, 0)
165
166
167def test_parse_flac_streaminfo_too_short() -> None:
168    """Header shorter than 42 bytes returns (0, 0)."""
169    assert KionMusicStreamingManager._parse_flac_streaminfo(b"fLaC\x00\x00") == (0, 0)
170
171
172def test_parse_mp4_audio_params_dfla() -> None:
173    """DfLa (FLAC-in-MP4) box is parsed via embedded STREAMINFO."""
174    result = KionMusicStreamingManager._parse_mp4_audio_params(_mp4_dfla_header(44100, 16))
175    assert result == (44100, 16)
176
177
178def test_parse_mp4_audio_params_mp4a_fallback() -> None:
179    """When dfLa is absent, mp4a AudioSampleEntry provides sample_rate/bit_depth."""
180    result = KionMusicStreamingManager._parse_mp4_audio_params(_mp4_mp4a_header(48000, 16))
181    assert result == (48000, 16)
182
183
184def test_parse_mp4_audio_params_no_box_returns_zero() -> None:
185    """Buffer containing neither dfLa nor mp4a returns (0, 0)."""
186    assert KionMusicStreamingManager._parse_mp4_audio_params(b"\x00" * 128) == (0, 0)
187
188
189def test_get_content_type_flac_mp4_returns_mp4_flac(
190    streaming_manager: KionMusicStreamingManager,
191) -> None:
192    """flac-mp4 → (MP4 container, FLAC codec), matching yandex_music convention."""
193    assert streaming_manager._get_content_type("flac-mp4") == (
194        ContentType.MP4,
195        ContentType.FLAC,
196    )
197    assert streaming_manager._get_content_type("FLAC-MP4") == (
198        ContentType.MP4,
199        ContentType.FLAC,
200    )
201
202
203def test_get_content_type_aac_mp4_returns_mp4_aac(
204    streaming_manager: KionMusicStreamingManager,
205) -> None:
206    """aac-mp4 / he-aac-mp4 → (MP4 container, AAC codec)."""
207    assert streaming_manager._get_content_type("aac-mp4") == (
208        ContentType.MP4,
209        ContentType.AAC,
210    )
211    assert streaming_manager._get_content_type("he-aac-mp4") == (
212        ContentType.MP4,
213        ContentType.AAC,
214    )
215
216
217def test_get_content_type_plain_codecs(
218    streaming_manager: KionMusicStreamingManager,
219) -> None:
220    """Plain codecs report themselves as content_type with UNKNOWN codec_type."""
221    assert streaming_manager._get_content_type("flac") == (ContentType.FLAC, ContentType.UNKNOWN)
222    assert streaming_manager._get_content_type("mp3") == (ContentType.MP3, ContentType.UNKNOWN)
223    assert streaming_manager._get_content_type("mpeg") == (ContentType.MP3, ContentType.UNKNOWN)
224    assert streaming_manager._get_content_type("aac") == (ContentType.AAC, ContentType.UNKNOWN)
225    assert streaming_manager._get_content_type("he-aac") == (ContentType.AAC, ContentType.UNKNOWN)
226
227
228def _make_stream_details(
229    decryption_key: str, url: str = "https://example.com/enc.flac"
230) -> StreamDetails:
231    """Build a minimal StreamDetails for get_audio_stream tests."""
232    return StreamDetails(
233        provider="kion_music_instance",
234        item_id="test_track",
235        audio_format=AudioFormat(content_type=ContentType.FLAC),
236        data={
237            "url": url,
238            "codec": "flac",
239            "transport": "encraw",
240            "bit_rate": 0,
241            "fi_quality": "lossless",
242            "fi_codecs": "flac-mp4,flac",
243            "decryption_key": decryption_key,
244        },
245    )
246
247
248async def test_get_audio_stream_retries_on_payload_error_then_raises(
249    streaming_manager: KionMusicStreamingManager,
250    monkeypatch: pytest.MonkeyPatch,
251) -> None:
252    """ClientPayloadError causes retries; raises MediaNotFoundError after max retries."""
253    get_audio_stream = getattr(streaming_manager, "get_audio_stream", None)
254    if get_audio_stream is None:
255        pytest.skip("get_audio_stream not available in this provider version")
256
257    async def _no_sleep(_: float) -> None:
258        pass
259
260    monkeypatch.setattr(asyncio, "sleep", _no_sleep)
261
262    class _DroppingContent:
263        async def iter_chunked(self, n: int) -> Any:
264            raise ClientPayloadError("Connection dropped")
265            yield b""  # type: ignore[unreachable]  # makes this an async generator
266
267    class _DroppingResponse:
268        status = 200
269        content = _DroppingContent()
270
271        def raise_for_status(self) -> None:
272            pass
273
274    class _DroppingContext:
275        async def __aenter__(self) -> _DroppingResponse:
276            return _DroppingResponse()
277
278        async def __aexit__(self, *args: object) -> None:
279            pass
280
281    class _FakeHttpSession:
282        def get(self, url: str, headers: Any = None, **kwargs: Any) -> _DroppingContext:
283            return _DroppingContext()
284
285    streaming_manager_mass: Any = streaming_manager.mass
286    streaming_manager_mass.http_session = _FakeHttpSession()
287    streamdetails = _make_stream_details(decryption_key="00" * 16)  # valid 16-byte AES key
288
289    with pytest.raises(MediaNotFoundError, match="retries were exhausted"):
290        async for _ in get_audio_stream(streamdetails):
291            pass
292