music-assistant-server

4.9 KBPY
test_direct_pcm_stream.py
4.9 KB154 lines • python
1"""Tests for the direct-PCM stream helper on the streams controller."""
2
3from __future__ import annotations
4
5from collections.abc import AsyncGenerator
6from typing import Any, cast
7from unittest.mock import MagicMock
8
9import pytest
10from music_assistant_models.enums import ContentType, MediaType, StreamType
11from music_assistant_models.media_items import AudioFormat
12from music_assistant_models.player import PlayerMedia
13from music_assistant_models.queue_item import QueueItem
14from music_assistant_models.streamdetails import StreamDetails
15
16from music_assistant.controllers.streams.controller import StreamsController
17
18PCM_FORMAT = AudioFormat(
19    content_type=ContentType.PCM_S16LE,
20    sample_rate=44100,
21    bit_depth=16,
22    channels=2,
23)
24
25QUEUE_ID = "player-1"
26QUEUE_ITEM_ID = "item-1"
27
28
29def _queue_item(seek_position: int) -> QueueItem:
30    """Build an audiobook queue item whose streamdetails carry the given seek position."""
31    return QueueItem(
32        queue_id=QUEUE_ID,
33        queue_item_id=QUEUE_ITEM_ID,
34        name="Some Audiobook",
35        duration=7200,
36        streamdetails=StreamDetails(
37            provider="builtin",
38            item_id="book-1",
39            audio_format=PCM_FORMAT,
40            media_type=MediaType.AUDIOBOOK,
41            stream_type=StreamType.HTTP,
42            path="http://example.com/book.mp3",
43            duration=7200,
44            can_seek=True,
45            allow_seek=True,
46            queue_id=QUEUE_ID,
47            seek_position=seek_position,
48        ),
49    )
50
51
52def _controller(queue_item: QueueItem) -> tuple[StreamsController, dict[str, Any]]:
53    """
54    Build a streams controller that records the kwargs of the single item stream call.
55
56    The queue itself is absent, which keeps the request on the non-flow (single item)
57    branch: crossfade and audio overlay both need a queue to force flow mode.
58
59    :param queue_item: The item the controller resolves the stream request to.
60    """
61    mass = MagicMock()
62    mass.config.get_raw_core_config_value.return_value = "GLOBAL"
63    mass.player_queues.get.return_value = None
64    mass.player_queues.get_item.return_value = queue_item
65    controller = StreamsController(mass)
66    call_kwargs: dict[str, Any] = {}
67
68    def _record(**kwargs: Any) -> object:
69        call_kwargs.update(kwargs)
70        return object()
71
72    controller.audio = MagicMock()
73    controller.audio.get_queue_item_stream.side_effect = _record
74    return controller, call_kwargs
75
76
77@pytest.mark.parametrize("seek_position", [1800, 0])
78def test_single_item_stream_forwards_the_seek_position(seek_position: int) -> None:
79    """A resumed (or seeked) item is streamed from its seek position, not from the start."""
80    queue_item = _queue_item(seek_position)
81    controller, call_kwargs = _controller(queue_item)
82
83    controller.get_stream(
84        PlayerMedia(
85            uri="library://audiobook/1",
86            media_type=MediaType.AUDIOBOOK,
87            source_id=QUEUE_ID,
88            queue_item_id=QUEUE_ITEM_ID,
89        ),
90        PCM_FORMAT,
91    )
92
93    assert call_kwargs["seek_position"] == seek_position
94
95
96async def _pcm_chunks() -> AsyncGenerator[bytes]:
97    """Yield two PCM chunks so the gauge can be observed mid-stream."""
98    yield b"chunk-1"
99    yield b"chunk-2"
100
101
102def _pcm_stream_controller() -> StreamsController:
103    """Build a controller whose single item stream yields real PCM chunks."""
104    controller, _ = _controller(_queue_item(0))
105    audio = cast("Any", controller.audio)
106    audio.get_queue_item_stream.side_effect = lambda **_kwargs: _pcm_chunks()
107    return controller
108
109
110@pytest.mark.asyncio
111async def test_direct_pcm_stream_counts_as_an_active_output_stream() -> None:
112    """A player consuming raw PCM registers as playing, so analysis yields CPU to it."""
113    controller = _pcm_stream_controller()
114
115    stream = controller.get_stream(
116        PlayerMedia(
117            uri="library://audiobook/1",
118            media_type=MediaType.AUDIOBOOK,
119            source_id=QUEUE_ID,
120            queue_item_id=QUEUE_ITEM_ID,
121        ),
122        PCM_FORMAT,
123    )
124    assert controller.output_stream_active() is False
125
126    chunks = []
127    async for chunk in stream:
128        chunks.append(chunk)
129        assert controller.output_stream_active() is True
130
131    assert chunks == [b"chunk-1", b"chunk-2"]
132    assert controller.output_stream_active() is False
133
134
135@pytest.mark.asyncio
136async def test_abandoned_direct_pcm_stream_releases_the_gauge() -> None:
137    """A consumer that stops mid-stream releases its count, so analysis regains its budget."""
138    controller = _pcm_stream_controller()
139
140    stream = controller.get_stream(
141        PlayerMedia(
142            uri="library://audiobook/1",
143            media_type=MediaType.AUDIOBOOK,
144            source_id=QUEUE_ID,
145            queue_item_id=QUEUE_ITEM_ID,
146        ),
147        PCM_FORMAT,
148    )
149    assert await anext(stream) == b"chunk-1"
150    assert controller.output_stream_active() is True
151
152    await stream.aclose()
153    assert controller.output_stream_active() is False
154