wehub-resource-sync
351 lines
9.8 KiB
Python
351 lines
9.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from livekit import rtc
|
|
from livekit.agents.voice.io import (
|
|
AgentOutput,
|
|
AudioOutput,
|
|
AudioOutputCapabilities,
|
|
PlaybackFinishedEvent,
|
|
_AudioSinkProxy,
|
|
)
|
|
|
|
from .fake_io import FakeAudioOutput
|
|
|
|
pytestmark = pytest.mark.unit
|
|
|
|
|
|
def _make_agent_output() -> AgentOutput:
|
|
return AgentOutput(lambda: None, lambda: None, lambda: None)
|
|
|
|
|
|
def _silence(duration_s: float = 0.01, sample_rate: int = 16000) -> rtc.AudioFrame:
|
|
n = int(sample_rate * duration_s)
|
|
return rtc.AudioFrame(
|
|
data=b"\x00\x00" * n,
|
|
sample_rate=sample_rate,
|
|
num_channels=1,
|
|
samples_per_channel=n,
|
|
)
|
|
|
|
|
|
class _PassthroughWrapper(AudioOutput):
|
|
"""Minimal wrapper that forwards capture_frame/flush/clear_buffer through
|
|
next_in_chain. Used to exercise the auto-wrap and swap mechanics without
|
|
pulling in TranscriptSynchronizer or RecorderAudioOutput.
|
|
"""
|
|
|
|
def __init__(self, *, next_in_chain: AudioOutput) -> None:
|
|
super().__init__(
|
|
label="Passthrough",
|
|
capabilities=AudioOutputCapabilities(pause=True),
|
|
next_in_chain=next_in_chain,
|
|
)
|
|
|
|
async def capture_frame(self, frame: rtc.AudioFrame) -> None:
|
|
await super().capture_frame(frame)
|
|
assert self.next_in_chain is not None
|
|
await self.next_in_chain.capture_frame(frame)
|
|
|
|
def flush(self) -> None:
|
|
super().flush()
|
|
assert self.next_in_chain is not None
|
|
self.next_in_chain.flush()
|
|
|
|
def clear_buffer(self) -> None:
|
|
assert self.next_in_chain is not None
|
|
self.next_in_chain.clear_buffer()
|
|
|
|
|
|
class _TrackingSink(AudioOutput):
|
|
"""Leaf sink that records attach/detach calls for assertion."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__(label="TrackingSink", capabilities=AudioOutputCapabilities(pause=True))
|
|
self.attached_calls = 0
|
|
self.detached_calls = 0
|
|
|
|
async def capture_frame(self, frame: rtc.AudioFrame) -> None:
|
|
await super().capture_frame(frame)
|
|
|
|
def flush(self) -> None:
|
|
super().flush()
|
|
|
|
def clear_buffer(self) -> None:
|
|
pass
|
|
|
|
def on_attached(self) -> None:
|
|
self.attached_calls += 1
|
|
super().on_attached()
|
|
|
|
def on_detached(self) -> None:
|
|
self.detached_calls += 1
|
|
super().on_detached()
|
|
|
|
|
|
# ---------- auto-wrap ----------
|
|
|
|
|
|
def test_auto_wrap_inserts_proxy_above_bare_leaf() -> None:
|
|
leaf = FakeAudioOutput()
|
|
wrapper = _PassthroughWrapper(next_in_chain=leaf)
|
|
|
|
assert isinstance(wrapper.next_in_chain, _AudioSinkProxy)
|
|
assert wrapper.next_in_chain.next_in_chain is leaf
|
|
|
|
|
|
def test_auto_wrap_skipped_when_passed_an_existing_proxy() -> None:
|
|
leaf = FakeAudioOutput()
|
|
proxy = _AudioSinkProxy(leaf)
|
|
wrapper = _PassthroughWrapper(next_in_chain=proxy)
|
|
|
|
assert wrapper.next_in_chain is proxy
|
|
|
|
|
|
def test_auto_wrap_skipped_when_wrapping_a_non_leaf() -> None:
|
|
leaf = FakeAudioOutput()
|
|
inner = _PassthroughWrapper(next_in_chain=leaf)
|
|
outer = _PassthroughWrapper(next_in_chain=inner)
|
|
|
|
# outer should reference inner directly — no proxy interposed between them
|
|
assert outer.next_in_chain is inner
|
|
|
|
|
|
# ---------- replace_audio_tail ----------
|
|
|
|
|
|
def test_replace_audio_tail_swaps_proxy_inner() -> None:
|
|
leaf_a = FakeAudioOutput()
|
|
leaf_b = FakeAudioOutput()
|
|
output = _make_agent_output()
|
|
wrapper = _PassthroughWrapper(next_in_chain=leaf_a)
|
|
output.audio = wrapper
|
|
|
|
output.replace_audio_tail(leaf_b)
|
|
|
|
# wrapper chain intact; only the leaf swapped
|
|
assert output.audio is wrapper
|
|
proxy = wrapper.next_in_chain
|
|
assert isinstance(proxy, _AudioSinkProxy)
|
|
assert proxy.next_in_chain is leaf_b
|
|
|
|
|
|
def test_replace_audio_tail_falls_back_when_no_proxy() -> None:
|
|
leaf = FakeAudioOutput()
|
|
output = _make_agent_output()
|
|
# no wrapper chain yet
|
|
|
|
output.replace_audio_tail(leaf)
|
|
|
|
assert output.audio is leaf
|
|
|
|
|
|
# ---------- proxy invariants ----------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_accepts_wrapper_chain_as_inner() -> None:
|
|
leaf = FakeAudioOutput()
|
|
wrapped_sink = _PassthroughWrapper(next_in_chain=leaf)
|
|
proxy = _AudioSinkProxy(wrapped_sink)
|
|
|
|
assert proxy.next_in_chain is wrapped_sink
|
|
|
|
# events from the inner leaf still bubble up through proxy
|
|
received: list[PlaybackFinishedEvent] = []
|
|
proxy.on("playback_finished", received.append)
|
|
|
|
await proxy.capture_frame(_silence())
|
|
leaf.on_playback_finished(playback_position=1.0, interrupted=False)
|
|
|
|
assert len(received) == 1
|
|
assert received[0].playback_position == 1.0
|
|
|
|
|
|
# ---------- swap routing ----------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_swap_routes_playback_events_from_new_leaf() -> None:
|
|
leaf_a = FakeAudioOutput()
|
|
leaf_b = FakeAudioOutput()
|
|
wrapper = _PassthroughWrapper(next_in_chain=leaf_a)
|
|
proxy = wrapper.next_in_chain
|
|
assert isinstance(proxy, _AudioSinkProxy)
|
|
|
|
received: list[PlaybackFinishedEvent] = []
|
|
wrapper.on("playback_finished", received.append)
|
|
|
|
proxy.set_next_in_chain(leaf_b)
|
|
# capture after the swap so leaf_b has a segment to mark as finished
|
|
await wrapper.capture_frame(_silence())
|
|
|
|
leaf_b.on_playback_finished(playback_position=1.0, interrupted=False)
|
|
assert len(received) == 1
|
|
assert received[0].playback_position == 1.0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_swap_disconnects_old_leaf() -> None:
|
|
leaf_a = FakeAudioOutput()
|
|
leaf_b = FakeAudioOutput()
|
|
wrapper = _PassthroughWrapper(next_in_chain=leaf_a)
|
|
proxy = wrapper.next_in_chain
|
|
assert isinstance(proxy, _AudioSinkProxy)
|
|
|
|
# give leaf_a a segment so its on_playback_finished would normally emit
|
|
await wrapper.capture_frame(_silence())
|
|
|
|
received: list[PlaybackFinishedEvent] = []
|
|
wrapper.on("playback_finished", received.append)
|
|
|
|
proxy.set_next_in_chain(leaf_b)
|
|
|
|
# leaf_a is detached: any event it fires must not reach the wrapper
|
|
leaf_a.on_playback_finished(playback_position=0.5, interrupted=False)
|
|
assert received == []
|
|
|
|
|
|
# ---------- swap with in-flight playback ----------
|
|
|
|
|
|
class _ClearCountingSink(FakeAudioOutput):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.clear_calls = 0
|
|
|
|
def clear_buffer(self) -> None:
|
|
self.clear_calls += 1
|
|
super().clear_buffer()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_swap_finishes_pending_segment_as_interrupted() -> None:
|
|
leaf_a = FakeAudioOutput()
|
|
leaf_b = FakeAudioOutput()
|
|
wrapper = _PassthroughWrapper(next_in_chain=leaf_a)
|
|
proxy = wrapper.next_in_chain
|
|
assert isinstance(proxy, _AudioSinkProxy)
|
|
|
|
# a flushed segment still playing out on leaf_a (frames are pushed faster than realtime)
|
|
await wrapper.capture_frame(_silence(duration_s=1.0))
|
|
wrapper.flush()
|
|
|
|
received: list[PlaybackFinishedEvent] = []
|
|
wrapper.on("playback_finished", received.append)
|
|
|
|
proxy.set_next_in_chain(leaf_b)
|
|
|
|
# the pending segment must be finished as interrupted so wait_for_playout() doesn't hang
|
|
ev = await asyncio.wait_for(wrapper.wait_for_playout(), timeout=0.5)
|
|
assert ev.interrupted is True
|
|
assert ev.playback_position == pytest.approx(1.0)
|
|
assert len(received) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_swap_clears_old_sink_with_inflight_audio() -> None:
|
|
leaf_a = _ClearCountingSink()
|
|
leaf_b = FakeAudioOutput()
|
|
wrapper = _PassthroughWrapper(next_in_chain=leaf_a)
|
|
proxy = wrapper.next_in_chain
|
|
assert isinstance(proxy, _AudioSinkProxy)
|
|
|
|
await wrapper.capture_frame(_silence(duration_s=1.0))
|
|
wrapper.flush()
|
|
|
|
proxy.set_next_in_chain(leaf_b)
|
|
|
|
assert leaf_a.clear_calls == 1
|
|
|
|
|
|
def test_idle_swap_does_not_clear_old_sink() -> None:
|
|
leaf_a = _ClearCountingSink()
|
|
leaf_b = FakeAudioOutput()
|
|
wrapper = _PassthroughWrapper(next_in_chain=leaf_a)
|
|
proxy = wrapper.next_in_chain
|
|
assert isinstance(proxy, _AudioSinkProxy)
|
|
|
|
proxy.set_next_in_chain(leaf_b)
|
|
|
|
assert leaf_a.clear_calls == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_swap_mid_capture_segment_finishes_on_new_leaf() -> None:
|
|
leaf_a = FakeAudioOutput()
|
|
leaf_b = FakeAudioOutput()
|
|
wrapper = _PassthroughWrapper(next_in_chain=leaf_a)
|
|
proxy = wrapper.next_in_chain
|
|
assert isinstance(proxy, _AudioSinkProxy)
|
|
|
|
received: list[PlaybackFinishedEvent] = []
|
|
wrapper.on("playback_finished", received.append)
|
|
|
|
# swap in the middle of a segment, before flush
|
|
await wrapper.capture_frame(_silence(duration_s=0.05))
|
|
proxy.set_next_in_chain(leaf_b)
|
|
|
|
# no synthesized event: the segment continues on leaf_b, which reports it
|
|
assert received == []
|
|
|
|
await wrapper.capture_frame(_silence(duration_s=0.05))
|
|
wrapper.flush()
|
|
|
|
ev = await asyncio.wait_for(wrapper.wait_for_playout(), timeout=1.0)
|
|
assert ev.interrupted is False
|
|
assert len(received) == 1
|
|
|
|
|
|
# ---------- attached state ----------
|
|
|
|
|
|
def test_proxy_propagates_on_attached_to_current_inner() -> None:
|
|
sink = _TrackingSink()
|
|
proxy = _AudioSinkProxy(sink)
|
|
|
|
proxy.on_attached()
|
|
|
|
assert sink.attached_calls == 1
|
|
assert sink.detached_calls == 0
|
|
|
|
|
|
def test_proxy_propagates_on_detached_to_current_inner() -> None:
|
|
sink = _TrackingSink()
|
|
proxy = _AudioSinkProxy(sink)
|
|
|
|
proxy.on_attached()
|
|
proxy.on_detached()
|
|
|
|
assert sink.attached_calls == 1
|
|
assert sink.detached_calls == 1
|
|
|
|
|
|
def test_swap_while_attached_attaches_new_and_detaches_old() -> None:
|
|
sink_a = _TrackingSink()
|
|
sink_b = _TrackingSink()
|
|
proxy = _AudioSinkProxy(sink_a)
|
|
proxy.on_attached()
|
|
|
|
assert sink_a.attached_calls == 1
|
|
|
|
proxy.set_next_in_chain(sink_b)
|
|
|
|
assert sink_a.detached_calls == 1
|
|
assert sink_b.attached_calls == 1
|
|
|
|
|
|
def test_swap_while_detached_does_not_fire_lifecycle_hooks() -> None:
|
|
sink_a = _TrackingSink()
|
|
sink_b = _TrackingSink()
|
|
proxy = _AudioSinkProxy(sink_a)
|
|
# never attached
|
|
|
|
proxy.set_next_in_chain(sink_b)
|
|
|
|
assert sink_a.detached_calls == 0
|
|
assert sink_b.attached_calls == 0
|