music-assistant-server

1"""Tests for the MilkDrop visualizer audio tap."""
2
3from __future__ import annotations
4
5import struct
6from unittest.mock import AsyncMock, Mock
7
8import numpy as np
9from music_assistant_models.media_items import AudioFormat, MediaItemPalette
10
11from music_assistant.models.audio_analysis import AudioAnalysisData
12from music_assistant.providers.milkdrop_visualizer.tap import (
13    WAVE_SAMPLES,
14    Tap,
15    TapManager,
16    TrackCursor,
17    ViewerQueue,
18    pack_wave_frame,
19    palette_payload,
20    pcm_to_mono,
21    server_now_us,
22)
23
24PCM_FORMAT = AudioFormat(sample_rate=44100, bit_depth=16, channels=2)
25
26
27def _stereo_pcm(mono_values: list[int], *, dangling_sample: bool = False) -> bytes:
28    """Build interleaved 16-bit stereo PCM from mono values, duplicated across L/R."""
29    samples: list[int] = []
30    for value in mono_values:
31        samples.extend((value, value))
32    if dangling_sample:
33        # One unpaired sample: the reshape would fail without the drop guard.
34        samples.append(0)
35    return struct.pack(f"<{len(samples)}h", *samples)
36
37
38def _manager() -> TapManager:
39    """Return a tap manager whose provider is inert."""
40    provider = Mock()
41    provider.logger.getChild.return_value = Mock()
42    manager = TapManager(provider)
43    manager._schedule_beats = Mock()  # type: ignore[method-assign]
44    return manager
45
46
47def _cursor(next_chunk: int = 0, anchor_us: int = 0) -> TrackCursor:
48    """Return a cursor positioned at the start of a chunk."""
49    return TrackCursor(
50        item_id="item-1",
51        anchor_us=anchor_us,
52        next_chunk=next_chunk,
53        carry=np.zeros(0, dtype=np.float32),
54        carry_media=float(next_chunk),
55    )
56
57
58def test_mono_fold_averages_channels() -> None:
59    """A stereo chunk folds to one sample per frame, scaled into -1.0..1.0."""
60    mono = pcm_to_mono(_stereo_pcm([0, 16384, -16384]), PCM_FORMAT)
61    assert mono.size == 3
62    assert mono[0] == 0.0
63    assert round(float(mono[1]), 3) == 0.5
64    assert round(float(mono[2]), 3) == -0.5
65
66
67def test_mono_fold_drops_a_dangling_sample() -> None:
68    """A truncated chunk loses its unpaired sample instead of failing the reshape."""
69    assert pcm_to_mono(_stereo_pcm([0, 0], dangling_sample=True), PCM_FORMAT).size == 2
70
71
72def test_mono_fold_reads_packed_24_bit() -> None:
73    """24-bit PCM has no numpy dtype, so its sign handling is worth pinning down."""
74    fmt = AudioFormat(sample_rate=44100, bit_depth=24, channels=1)
75    # 0, +full scale - 1, -full scale
76    data = b"\x00\x00\x00" + b"\xff\xff\x7f" + b"\x00\x00\x80"
77    mono = pcm_to_mono(data, fmt)
78    assert mono[0] == 0.0
79    assert round(float(mono[1]), 3) == 1.0
80    assert round(float(mono[2]), 3) == -1.0
81
82
83def test_emits_one_frame_per_1024_samples() -> None:
84    """A one-second chunk yields a frame per full window, keeping the remainder back."""
85    manager = _manager()
86    tap = Tap("player-1")
87    cursor = _cursor()
88    manager._emit_chunk(tap, cursor, _stereo_pcm([0] * 44100), PCM_FORMAT)
89    assert len(tap.ring) == 44100 // WAVE_SAMPLES
90    assert cursor.carry.size == 44100 % WAVE_SAMPLES
91    assert cursor.next_chunk == 1
92
93
94def test_frame_is_stamped_at_the_end_of_its_window() -> None:
95    """A frame plays out at the anchor plus the media time its last sample sits at."""
96    manager = _manager()
97    tap = Tap("player-1")
98    cursor = _cursor(next_chunk=10, anchor_us=1_000_000)
99    manager._emit_chunk(tap, cursor, _stereo_pcm([0] * WAVE_SAMPLES), PCM_FORMAT)
100    tag, timestamp_us = struct.unpack(">Bq", tap.ring[0][:9])
101    assert tag == 22
102    assert len(tap.ring[0]) == 9 + WAVE_SAMPLES
103    # chunk 10 is media second 10, plus one 1024-sample window
104    expected_media = 10 + WAVE_SAMPLES / 44100
105    assert timestamp_us == 1_000_000 + int(expected_media * 1_000_000)
106
107
108def test_carry_continues_into_the_next_chunk() -> None:
109    """Samples left over from a chunk complete the first window of the next one."""
110    manager = _manager()
111    tap = Tap("player-1")
112    cursor = _cursor()
113    second = _stereo_pcm([0] * 44100)
114    manager._emit_chunk(tap, cursor, second, PCM_FORMAT)
115    manager._emit_chunk(tap, cursor, second, PCM_FORMAT)
116    # windows tile the two seconds end to end, rather than restarting per chunk
117    assert len(tap.ring) == (2 * 44100) // WAVE_SAMPLES
118    _, timestamp_us = struct.unpack(">Bq", tap.ring[-1][:9])
119    assert timestamp_us == int(len(tap.ring) * WAVE_SAMPLES / 44100 * 1_000_000)
120
121
122def test_carry_is_dropped_when_the_next_chunk_is_elsewhere() -> None:
123    """After a resync the leftover belongs to audio we are no longer continuing from."""
124    manager = _manager()
125    tap = Tap("player-1")
126    cursor = _cursor()
127    manager._emit_chunk(tap, cursor, _stereo_pcm([0] * 44100), PCM_FORMAT)
128    carried = cursor.carry.size
129    assert carried
130    cursor.next_chunk = 60
131    manager._emit_chunk(tap, cursor, _stereo_pcm([0] * 44100), PCM_FORMAT)
132    assert cursor.carry_media > 60
133
134
135def test_quantized_samples_are_offset_binary() -> None:
136    """Silence sits at 0x80, so a viewer reads the tail without knowing the scale."""
137    manager = _manager()
138    tap = Tap("player-1")
139    manager._emit_chunk(tap, _cursor(), _stereo_pcm([0] * WAVE_SAMPLES), PCM_FORMAT)
140    assert set(tap.ring[0][9:]) == {0x80}
141
142
143def test_align_keeps_a_cursor_that_still_matches() -> None:
144    """A cursor on the same track, in step with the queue, is left alone."""
145    manager = _manager()
146    tap = Tap("player-1")
147    item = Mock(queue_item_id="item-1")
148    buffer = Mock(first_buffered_chunk=0)
149    cursor = manager._align(tap, None, item, 5.0, buffer)
150    assert manager._align(tap, cursor, item, 5.0, buffer) is cursor
151
152
153def test_align_re_anchors_on_a_track_change() -> None:
154    """A new queue item drops what was scheduled from the old track's timeline."""
155    manager = _manager()
156    tap = Tap("player-1")
157    buffer = Mock(first_buffered_chunk=0)
158    cursor = manager._align(tap, None, Mock(queue_item_id="item-1"), 30.0, buffer)
159    tap.ring.append(b"stale")
160    queued = ViewerQueue()
161    tap.queues.add(queued)
162    new_cursor = manager._align(tap, cursor, Mock(queue_item_id="item-2"), 0.0, buffer)
163    assert new_cursor is not cursor
164    assert new_cursor.next_chunk == 0
165    assert not tap.ring
166    assert queued._items[0] == '{"type": "stream/clear"}'
167
168
169def test_align_re_anchors_on_a_seek() -> None:
170    """A playhead that jumps away from the anchored timeline restarts the cursor."""
171    manager = _manager()
172    tap = Tap("player-1")
173    item = Mock(queue_item_id="item-1")
174    buffer = Mock(first_buffered_chunk=0)
175    cursor = manager._align(tap, None, item, 5.0, buffer)
176    new_cursor = manager._align(tap, cursor, item, 120.0, buffer)
177    assert new_cursor is not cursor
178    assert new_cursor.next_chunk == 120
179
180
181def test_align_starts_inside_the_retained_window() -> None:
182    """A rolling buffer that has discarded the playhead is picked up where it starts."""
183    manager = _manager()
184    tap = Tap("player-1")
185    item = Mock(queue_item_id="item-1")
186    cursor = manager._align(tap, None, item, 5.0, Mock(first_buffered_chunk=90))
187    assert cursor.next_chunk == 90
188
189
190def test_playhead_tracks_playback_speed() -> None:
191    """At 2x, media time advances two seconds per wall-clock second from the anchor."""
192    cursor = _cursor(anchor_us=server_now_us() - 10_000_000)
193    cursor.speed = 2.0
194    assert abs(cursor.playhead() - 20.0) < 0.1
195    # and the inverse mapping stamps media second 20 at (roughly) now
196    assert abs(cursor.media_to_clock_us(20.0) - server_now_us()) < 100_000
197
198
199def test_align_keeps_a_speed_aware_cursor_in_step() -> None:
200    """A cursor anchored at 2x stays matched against a queue advancing in media-time."""
201    manager = _manager()
202    tap = Tap("player-1")
203    item = Mock(queue_item_id="item-1")
204    buffer = Mock(first_buffered_chunk=0)
205    cursor = manager._align(tap, None, item, 10.0, buffer, 2.0)
206    assert cursor.speed == 2.0
207    assert manager._align(tap, cursor, item, 10.0, buffer, 2.0) is cursor
208
209
210def test_align_scales_the_resync_threshold_by_speed() -> None:
211    """At 2x, report jitter inflates by the speed factor, so the threshold grows with it."""
212    manager = _manager()
213    tap = Tap("player-1")
214    item = Mock(queue_item_id="item-1")
215    buffer = Mock(first_buffered_chunk=0)
216    cursor = manager._align(tap, None, item, 10.0, buffer, 2.0)
217    # a 5s media-time gap is within the scaled 6s threshold, not a seek
218    assert manager._align(tap, cursor, item, 15.0, buffer, 2.0) is cursor
219    # beyond the scaled threshold it is a seek and re-anchors
220    assert manager._align(tap, cursor, item, 17.0, buffer, 2.0) is not cursor
221
222
223def test_align_re_anchors_on_a_speed_change() -> None:
224    """A playback speed change remaps media time to the clock, so the cursor restarts."""
225    manager = _manager()
226    tap = Tap("player-1")
227    item = Mock(queue_item_id="item-1")
228    buffer = Mock(first_buffered_chunk=0)
229    cursor = manager._align(tap, None, item, 10.0, buffer)
230    new_cursor = manager._align(tap, cursor, item, 10.0, buffer, 1.5)
231    assert new_cursor is not cursor
232    assert new_cursor.speed == 1.5
233
234
235def _beats_manager() -> TapManager:
236    """Return a tap manager with the real beat scheduling in place."""
237    provider = Mock()
238    provider.logger.getChild.return_value = Mock()
239    return TapManager(provider)
240
241
242def test_schedule_beats_rebuilds_from_cached_analysis() -> None:
243    """A re-anchor of an item whose analysis is cached reschedules in place, without a task."""
244    manager = _beats_manager()
245    tap = Tap("player-1")
246    tap.beats_analysis = ("item-1", AudioAnalysisData(beats=[1.0, 2.0], downbeats=[1.0]))
247    anchor_us = server_now_us()
248    manager._schedule_beats(tap, Mock(queue_item_id="item-1"), anchor_us)
249    manager.mass.create_task.assert_not_called()  # type: ignore[attr-defined]
250    assert [timestamp_us for timestamp_us, _ in tap.beats] == [
251        anchor_us + 1_000_000,
252        anchor_us + 2_000_000,
253    ]
254    # the downbeat flag survives the rebuild
255    assert tap.beats[0][1][9] == 1
256    assert tap.beats[1][1][9] == 0
257
258
259def test_fan_out_beats_scales_media_time_by_speed() -> None:
260    """At 2x a beat at media second 2 sounds one wall-clock second after the anchor."""
261    manager = _beats_manager()
262    tap = Tap("player-1")
263    anchor_us = server_now_us()
264    manager._fan_out_beats(tap, AudioAnalysisData(beats=[2.0]), anchor_us, 2.0)
265    assert [timestamp_us for timestamp_us, _ in tap.beats] == [anchor_us + 1_000_000]
266
267
268def test_reset_cancels_an_in_flight_beat_hydration() -> None:
269    """A timeline reset stops a pending hydration from landing beats for a dead track."""
270    tap = Tap("player-1")
271    task = Mock()
272    tap.beats_task = task
273    tap.reset('{"type": "stream/end"}')
274    task.cancel.assert_called_once()
275    assert tap.beats_task is None
276
277
278def test_schedule_beats_does_not_serve_another_item_from_cache() -> None:
279    """A cached analysis belongs to one item; any other item hydrates freshly."""
280    manager = _beats_manager()
281    tap = Tap("player-1")
282    tap.beats_analysis = ("item-1", AudioAnalysisData(beats=[1.0]))
283    manager._schedule_beats(tap, Mock(queue_item_id="item-2"), server_now_us())
284    manager.mass.create_task.assert_called_once()  # type: ignore[attr-defined]
285    assert not tap.beats
286
287
288async def test_hydrate_beats_caches_the_fetched_analysis() -> None:
289    """The fetched analysis is kept on the tap so the next re-anchor skips the query."""
290    manager = _beats_manager()
291    analysis = AudioAnalysisData(beats=[1.0])
292    manager.mass.streams.audio_analysis.get_audio_analysis = AsyncMock(  # type: ignore[method-assign]
293        return_value=analysis
294    )
295    tap = Tap("player-1")
296    await manager._hydrate_beats(tap, Mock(queue_item_id="item-1"), server_now_us())
297    assert tap.beats_analysis == ("item-1", analysis)
298    assert len(tap.beats) == 1
299
300
301def test_ring_with_only_future_frames_is_reported_stale() -> None:
302    """A ring whose oldest frame is ahead of now has nothing a fresh viewer can draw."""
303    tap = Tap("player-1")
304    assert not tap.has_only_future_frames()
305    tap.ring.append(pack_wave_frame(server_now_us() - 1_000_000, b"\x80" * WAVE_SAMPLES))
306    assert not tap.has_only_future_frames()
307    tap.ring.clear()
308    tap.ring.append(pack_wave_frame(server_now_us() + 60_000_000, b"\x80" * WAVE_SAMPLES))
309    assert tap.has_only_future_frames()
310
311
312async def test_read_once_realigns_when_requested() -> None:
313    """A requested realign drops a pinned-ahead cursor and restarts at the playhead."""
314    manager = _manager()
315    manager.provider.config.get_value.return_value = False  # type: ignore[attr-defined]
316    tap = Tap("player-1")
317    queue = Mock(corrected_elapsed_time=100.0, playback_speed=1.0)
318    item = Mock(queue_item_id="item-1")
319    buffer = Mock(first_buffered_chunk=0, pcm_format=PCM_FORMAT)
320    buffer.read_chunk_for_analysis = AsyncMock(return_value=_stereo_pcm([0] * 44100))
321    manager._playing_source = Mock(return_value=(queue, item, buffer))  # type: ignore[method-assign]
322    # a cursor pinned at the eviction edge but otherwise in sync with the queue
323    pinned = _cursor(next_chunk=200, anchor_us=server_now_us() - 100_000_000)
324    tap.realign_requested = True
325    cursor = await manager._read_once(tap, pinned)
326    assert cursor is not None
327    assert cursor is not pinned
328    assert cursor.next_chunk == 101
329    assert not tap.realign_requested
330
331
332async def test_read_once_ignores_realign_when_playhead_chunk_is_evicted() -> None:
333    """A realign past the eviction edge is dropped: healthy viewers keep their frames."""
334    manager = _manager()
335    manager.provider.config.get_value.return_value = False  # type: ignore[attr-defined]
336    tap = Tap("player-1")
337    tap.ring.append(b"frame")
338    queue = Mock(corrected_elapsed_time=100.0, playback_speed=1.0)
339    item = Mock(queue_item_id="item-1")
340    buffer = Mock(first_buffered_chunk=200, pcm_format=PCM_FORMAT)
341    buffer.read_chunk_for_analysis = AsyncMock(return_value=_stereo_pcm([0] * 44100))
342    manager._playing_source = Mock(return_value=(queue, item, buffer))  # type: ignore[method-assign]
343    # a cursor pinned at the eviction edge but otherwise in sync with the queue
344    pinned = _cursor(next_chunk=200, anchor_us=server_now_us() - 100_000_000)
345    tap.realign_requested = True
346    cursor = await manager._read_once(tap, pinned)
347    assert cursor is pinned
348    assert b"frame" in tap.ring
349    assert not tap.realign_requested
350
351
352def test_palette_payload_maps_every_field() -> None:
353    """A palette becomes the color@v1 payload the wire format documents."""
354    payload = palette_payload(MediaItemPalette(primary=(1, 2, 3)))
355    assert payload["primary"] == [1, 2, 3]
356    assert payload["accent"] is None
357
358
359def test_palette_payload_nulls_everything_without_a_palette() -> None:
360    """A track with no palette clears the previous track's tint."""
361    payload = palette_payload(None)
362    assert payload
363    assert all(value is None for value in payload.values())
364
365
366def test_viewer_queue_evicts_oldest_binary_frame_when_full() -> None:
367    """A stalled viewer loses waveform frames rather than stalling the tap."""
368    queue = ViewerQueue(capacity=2)
369    queue.push(b"first")
370    queue.push(b"second")
371    queue.push('{"type": "stream/clear"}')
372    drained = [queue._items[index] for index in range(len(queue._items))]
373    assert drained == [b"second", '{"type": "stream/clear"}']
374
375
376def test_viewer_queue_evicts_control_only_when_no_binary_left() -> None:
377    """Control messages are kept while any waveform frame can be dropped instead."""
378    queue = ViewerQueue(capacity=2)
379    queue.push('{"type": "stream/start"}')
380    queue.push('{"type": "stream/clear"}')
381    queue.push(b"frame")
382    drained = [queue._items[index] for index in range(len(queue._items))]
383    assert len(drained) == 2
384    assert '{"type": "stream/clear"}' in drained
385