wehub-resource-sync
96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pybase64 as base64
|
|
import pytest
|
|
|
|
from vllm.multimodal.media import AudioMediaIO
|
|
from vllm.multimodal.media.audio import load_audio
|
|
|
|
from ...conftest import AudioTestAssets
|
|
|
|
pytestmark = pytest.mark.cpu_test
|
|
|
|
ASSETS_DIR = Path(__file__).parent.parent / "assets"
|
|
assert ASSETS_DIR.exists()
|
|
|
|
|
|
@pytest.fixture
|
|
def dummy_audio():
|
|
return np.array([0.0, 0.1, 0.2, 0.3, 0.4], dtype=float)
|
|
|
|
|
|
@pytest.fixture
|
|
def dummy_audio_bytes(audio_assets: AudioTestAssets):
|
|
with open(audio_assets[0].get_local_path(), "rb") as f:
|
|
return f.read()
|
|
|
|
|
|
def test_audio_media_io_load_bytes(dummy_audio_bytes):
|
|
audio_io = AudioMediaIO()
|
|
out = audio_io.load_bytes(dummy_audio_bytes)
|
|
assert isinstance(out[0], np.ndarray)
|
|
assert out[1] == 16000
|
|
|
|
|
|
def test_audio_media_io_load_base64(dummy_audio_bytes):
|
|
audio_io = AudioMediaIO()
|
|
encoded = base64.b64encode(dummy_audio_bytes).decode("utf-8")
|
|
out = audio_io.load_base64("audio/wav", encoded)
|
|
assert isinstance(out[0], np.ndarray)
|
|
assert out[1] == 16000
|
|
|
|
|
|
def test_audio_media_io_load_file(audio_assets: AudioTestAssets):
|
|
audio_io = AudioMediaIO()
|
|
path = audio_assets[0].get_local_path()
|
|
out = audio_io.load_file(path)
|
|
assert isinstance(out[0], np.ndarray)
|
|
assert out[1] == 16000
|
|
|
|
|
|
def test_audio_media_io_encode_base64(dummy_audio):
|
|
audio_io = AudioMediaIO()
|
|
media = (dummy_audio, 16000)
|
|
with patch("soundfile.write") as mock_write:
|
|
|
|
def write_to_buffer(buffer, *_args, **_kwargs):
|
|
buffer.write(b"dummy_wav_data")
|
|
|
|
mock_write.side_effect = write_to_buffer
|
|
|
|
out = audio_io.encode_base64(media)
|
|
decoded = base64.b64decode(out)
|
|
assert decoded == b"dummy_wav_data"
|
|
mock_write.assert_called_once()
|
|
|
|
|
|
def test_load_audio_max_duration_respected(dummy_audio_bytes):
|
|
"""Valid audio within the duration limit should load successfully."""
|
|
from io import BytesIO
|
|
|
|
y, sr = load_audio(BytesIO(dummy_audio_bytes), sr=None, max_duration_s=3600)
|
|
assert isinstance(y, np.ndarray)
|
|
assert len(y) > 0
|
|
|
|
|
|
def test_load_audio_max_duration_rejected(dummy_audio_bytes):
|
|
"""Audio exceeding the duration limit must be rejected during decode."""
|
|
from io import BytesIO
|
|
|
|
with pytest.raises(ValueError, match="exceeds maximum allowed duration"):
|
|
load_audio(BytesIO(dummy_audio_bytes), sr=None, max_duration_s=0.0001)
|
|
|
|
|
|
def test_audio_media_io_from_video(video_assets):
|
|
audio_io = AudioMediaIO()
|
|
video_path = video_assets[0].video_path
|
|
with open(video_path, "rb") as f:
|
|
audio, sr = audio_io.load_bytes(f.read())
|
|
audio_ref, sr_ref = load_audio(video_path, sr=None)
|
|
assert sr == sr_ref
|
|
np.testing.assert_allclose(audio_ref, audio, atol=1e-4)
|