65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import sys
|
|
import types
|
|
import wave
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
from app.livekit.runtime.wav_audio import wav_audio_frames, wav_duration_ms
|
|
|
|
|
|
class _AudioFrame:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
data: bytes,
|
|
sample_rate: int,
|
|
num_channels: int,
|
|
samples_per_channel: int,
|
|
) -> None:
|
|
self.data = data
|
|
self.sample_rate = sample_rate
|
|
self.num_channels = num_channels
|
|
self.samples_per_channel = samples_per_channel
|
|
|
|
|
|
def _write_wav(path: Path, *, samples: int, sample_rate: int = 1000) -> None:
|
|
with wave.open(str(path), "wb") as wav:
|
|
wav.setnchannels(1)
|
|
wav.setsampwidth(2)
|
|
wav.setframerate(sample_rate)
|
|
wav.writeframes(b"\x01\x02" * samples)
|
|
|
|
|
|
def test_wav_audio_frames_pads_last_frame_and_adds_tail_silence(tmp_path: Path) -> None:
|
|
wav_path = tmp_path / "audio.wav"
|
|
_write_wav(wav_path, samples=25)
|
|
|
|
livekit_module = types.ModuleType("livekit")
|
|
rtc_module = types.ModuleType("livekit.rtc")
|
|
rtc_module.AudioFrame = _AudioFrame
|
|
livekit_module.rtc = rtc_module
|
|
|
|
async def _collect():
|
|
with mock.patch.dict(sys.modules, {"livekit": livekit_module, "livekit.rtc": rtc_module}):
|
|
return [
|
|
frame
|
|
async for frame in wav_audio_frames(
|
|
str(wav_path),
|
|
frame_duration_ms=20,
|
|
tail_silence_ms=40,
|
|
)
|
|
]
|
|
|
|
frames = asyncio.run(_collect())
|
|
|
|
assert len(frames) == 4
|
|
assert [frame.samples_per_channel for frame in frames] == [20, 20, 20, 20]
|
|
assert frames[1].data[:10] == b"\x01\x02" * 5
|
|
assert frames[1].data[10:] == b"\x00" * 30
|
|
assert frames[2].data == b"\x00" * 40
|
|
assert frames[3].data == b"\x00" * 40
|
|
assert wav_duration_ms(str(wav_path)) == 25
|