wehub-resource-sync
268 lines
8.2 KiB
Python
268 lines
8.2 KiB
Python
"""Tests for AudioEmitter: is_final, timed transcripts, duration tracking, segments."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import io
|
|
import wave
|
|
|
|
import pytest
|
|
|
|
from livekit.agents import tts, utils
|
|
from livekit.agents.types import USERDATA_TIMED_TRANSCRIPT, TimedString
|
|
|
|
pytestmark = [pytest.mark.unit, pytest.mark.concurrent]
|
|
|
|
|
|
def _make_pcm(sample_rate: int, num_channels: int, duration_ms: int) -> bytes:
|
|
num_samples = sample_rate * duration_ms // 1000 * num_channels
|
|
return b"\x00\x00" * num_samples
|
|
|
|
|
|
def _make_wav(sample_rate: int, num_channels: int, duration_ms: int) -> bytes:
|
|
num_samples = sample_rate * duration_ms // 1000
|
|
buf = io.BytesIO()
|
|
with wave.open(buf, "wb") as w:
|
|
w.setnchannels(num_channels)
|
|
w.setsampwidth(2)
|
|
w.setframerate(sample_rate)
|
|
w.writeframes(b"\x00\x00" * (num_samples * num_channels))
|
|
return buf.getvalue()
|
|
|
|
|
|
SR = 24000
|
|
NC = 1
|
|
|
|
|
|
async def _run_emitter(
|
|
produce_fn, *, stream: bool = False, expect_finals: int = 1, mime_type: str = "audio/pcm"
|
|
):
|
|
"""Run emitter, collect events until we've seen expect_finals is_final events."""
|
|
dst_ch = utils.aio.Chan[tts.SynthesizedAudio]()
|
|
emitter = tts.AudioEmitter(label="test", dst_ch=dst_ch)
|
|
emitter.initialize(
|
|
request_id="req1",
|
|
sample_rate=SR,
|
|
num_channels=NC,
|
|
mime_type=mime_type,
|
|
stream=stream,
|
|
)
|
|
|
|
events: list[tts.SynthesizedAudio] = []
|
|
finals_seen = 0
|
|
|
|
async def collect():
|
|
nonlocal finals_seen
|
|
async for ev in dst_ch:
|
|
events.append(ev)
|
|
if ev.is_final:
|
|
finals_seen += 1
|
|
if finals_seen >= expect_finals:
|
|
return
|
|
|
|
ct = asyncio.create_task(collect())
|
|
produce_fn(emitter)
|
|
await emitter.join()
|
|
await ct
|
|
return emitter, events
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_basic_is_final():
|
|
def produce(e):
|
|
for _ in range(50):
|
|
e.push(_make_pcm(SR, NC, 10))
|
|
e.end_input()
|
|
|
|
_, events = await _run_emitter(produce)
|
|
|
|
assert len(events) > 0
|
|
assert events[-1].is_final
|
|
assert all(not e.is_final for e in events[:-1])
|
|
total = sum(e.frame.duration for e in events)
|
|
assert abs(total - 0.5) < 0.02
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_single_tiny_push():
|
|
def produce(e):
|
|
e.push(_make_pcm(SR, NC, 5))
|
|
e.end_input()
|
|
|
|
_, events = await _run_emitter(produce)
|
|
assert len(events) >= 1
|
|
assert events[-1].is_final
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_timed_transcripts_all_preserved():
|
|
words = ["hello", "world", "this", "is", "a", "test"]
|
|
|
|
def produce(e):
|
|
for i, word in enumerate(words):
|
|
e.push(_make_pcm(SR, NC, 80))
|
|
e.push_timed_transcript(
|
|
TimedString(text=word, start_time=i * 0.08, end_time=(i + 1) * 0.08)
|
|
)
|
|
e.end_input()
|
|
|
|
_, events = await _run_emitter(produce)
|
|
|
|
recovered = []
|
|
for ev in events:
|
|
for t in ev.frame.userdata.get(USERDATA_TIMED_TRANSCRIPT, []):
|
|
if isinstance(t, TimedString):
|
|
recovered.append(str(t))
|
|
|
|
assert recovered == words, f"expected {words}, got {recovered}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_timed_transcripts_have_timestamps():
|
|
def produce(e):
|
|
e.push(_make_pcm(SR, NC, 200))
|
|
e.push_timed_transcript(TimedString(text="hello", start_time=0.0, end_time=0.1))
|
|
e.push_timed_transcript(TimedString(text="world", start_time=0.1, end_time=0.2))
|
|
e.end_input()
|
|
|
|
_, events = await _run_emitter(produce)
|
|
|
|
all_ts: list[TimedString] = []
|
|
for ev in events:
|
|
for t in ev.frame.userdata.get(USERDATA_TIMED_TRANSCRIPT, []):
|
|
if isinstance(t, TimedString):
|
|
all_ts.append(t)
|
|
|
|
assert len(all_ts) == 2
|
|
assert str(all_ts[0]) == "hello"
|
|
assert all_ts[0].start_time == 0.0
|
|
assert str(all_ts[1]) == "world"
|
|
assert all_ts[1].start_time == 0.1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_segments_each_have_is_final():
|
|
def produce(e):
|
|
for seg_idx in range(3):
|
|
e.start_segment(segment_id=f"seg_{seg_idx}")
|
|
for _ in range(10):
|
|
e.push(_make_pcm(SR, NC, 10))
|
|
e.end_segment()
|
|
e.end_input()
|
|
|
|
_, events = await _run_emitter(produce, stream=True, expect_finals=3)
|
|
|
|
finals = [e for e in events if e.is_final]
|
|
assert len(finals) == 3
|
|
total = sum(e.frame.duration for e in events)
|
|
assert abs(total - 0.3) < 0.02
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_transcripts_across_segments():
|
|
def produce(e):
|
|
e.start_segment(segment_id="seg_0")
|
|
e.push(_make_pcm(SR, NC, 200))
|
|
e.push_timed_transcript(TimedString(text="first", start_time=0.0, end_time=0.2))
|
|
e.end_segment()
|
|
|
|
e.start_segment(segment_id="seg_1")
|
|
e.push(_make_pcm(SR, NC, 200))
|
|
e.push_timed_transcript(TimedString(text="second", start_time=0.0, end_time=0.2))
|
|
e.end_segment()
|
|
|
|
e.end_input()
|
|
|
|
_, events = await _run_emitter(produce, stream=True, expect_finals=2)
|
|
|
|
all_words = []
|
|
for ev in events:
|
|
for t in ev.frame.userdata.get(USERDATA_TIMED_TRANSCRIPT, []):
|
|
if isinstance(t, TimedString):
|
|
all_words.append(str(t))
|
|
|
|
assert "first" in all_words
|
|
assert "second" in all_words
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pushed_duration_accurate():
|
|
def produce(e):
|
|
e.push(_make_pcm(SR, NC, 500))
|
|
e.end_input()
|
|
|
|
emitter, _ = await _run_emitter(produce)
|
|
assert abs(emitter.pushed_duration() - 0.5) < 0.02
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_marker_does_not_inflate_duration():
|
|
"""flush() followed by end_input() with no new audio: marker should not inflate duration."""
|
|
|
|
def produce(e):
|
|
e.push(_make_pcm(SR, NC, 100))
|
|
e.flush()
|
|
e.end_input()
|
|
|
|
emitter, events = await _run_emitter(produce)
|
|
assert events[-1].is_final
|
|
assert abs(emitter.pushed_duration() - 0.1) < 0.02
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flush_mid_stream_preserves_decoder_state():
|
|
"""Regression: flush() between two chunks of the same encoded stream must not
|
|
tear the decoder down.
|
|
|
|
A stateful codec (WAV/OGG/MP3) parses headers once; if the FlushSegment branch
|
|
in the decoder path discards the decoder, the next chunk — which is mid-file
|
|
payload (raw PCM continuation, no RIFF) — gets fed to a fresh decoder that
|
|
expects a header and fails. The user-visible symptom is
|
|
``ValueError: Invalid WAV file: missing RIFF/WAVE: ...`` and silent audio loss.
|
|
|
|
flush() must release the held-back tail without touching the decoder.
|
|
"""
|
|
duration_ms = 500
|
|
wav = _make_wav(SR, NC, duration_ms)
|
|
# Split somewhere inside the data chunk so chunk2 has no RIFF header — this is
|
|
# the exact shape the streaming decoder branch must handle.
|
|
split = 100
|
|
chunk1, chunk2 = wav[:split], wav[split:]
|
|
assert chunk1[:4] == b"RIFF"
|
|
assert chunk2[:4] != b"RIFF"
|
|
|
|
def produce(e):
|
|
e.push(chunk1)
|
|
e.flush()
|
|
e.push(chunk2)
|
|
e.end_input()
|
|
|
|
emitter, events = await _run_emitter(produce, mime_type="audio/wav")
|
|
|
|
expected_samples = SR * duration_ms // 1000
|
|
total_samples = sum(ev.frame.samples_per_channel for ev in events)
|
|
# Accept a small tolerance for the trailing-tail handling, but anything close
|
|
# to the full payload means the decoder survived the flush. Pre-fix this
|
|
# would drop nearly everything (a few hundred samples of pre-flush audio).
|
|
assert total_samples >= int(expected_samples * 0.95), (
|
|
f"lost samples after mid-stream flush: got {total_samples}, expected ~{expected_samples}"
|
|
)
|
|
assert events[-1].is_final
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_every_frame_has_timed_transcript_metadata():
|
|
"""Every emitted frame must have the USERDATA_TIMED_TRANSCRIPT key in userdata."""
|
|
|
|
def produce(e):
|
|
e.push(_make_pcm(SR, NC, 300))
|
|
e.push_timed_transcript(TimedString(text="word", start_time=0.0, end_time=0.3))
|
|
e.end_input()
|
|
|
|
_, events = await _run_emitter(produce)
|
|
|
|
for i, ev in enumerate(events):
|
|
assert USERDATA_TIMED_TRANSCRIPT in ev.frame.userdata, (
|
|
f"frame {i} (is_final={ev.is_final}) missing {USERDATA_TIMED_TRANSCRIPT} in userdata"
|
|
)
|