"""Integration tests for wait_for_participant, wait_for_track_publication, and wait_for_agent. These tests require a running LiveKit server (bootstrapped via the livekit_server fixture). """ from __future__ import annotations import asyncio import contextlib import os import uuid from collections.abc import Awaitable, Callable from contextlib import asynccontextmanager import pytest from livekit import api, rtc from livekit.agents.utils.participant import ( wait_for_agent, wait_for_participant, wait_for_track_publication, ) from livekit.agents.voice.room_io._output import _ParticipantAudioOutput from .lk_server import LK_API_KEY, LK_API_SECRET, LK_URL, livekit_server # noqa: F401 from .utils.audio_test import AudioEnergyMonitor, SineToneSource from .utils.livekit_test import ( connect_room as _connect_e2e_room, make_room_name as _e2e_room_name, simulate_full_reconnect, simulate_resume, wait_for_event, ) pytestmark = [pytest.mark.unit, pytest.mark.concurrent] TIMEOUT = 5.0 def _make_token( identity: str, room: str, *, kind: str | None = None, agent: bool = False, ) -> str: token = ( api.AccessToken(LK_API_KEY, LK_API_SECRET) .with_identity(identity) .with_grants( api.VideoGrants( room_join=True, room=room, can_publish=True, can_subscribe=True, agent=agent, ) ) ) if kind is not None: token = token.with_kind(kind) return token.to_jwt() @asynccontextmanager async def connect_room( identity: str, room_name: str, *, kind: str | None = None, agent: bool = False ): room = rtc.Room() token = _make_token(identity, room_name, kind=kind, agent=agent) await room.connect(LK_URL, token) try: yield room finally: await room.disconnect() def _room_name() -> str: return f"test-{uuid.uuid4()}" async def _publish_audio_track(room: rtc.Room) -> None: source = rtc.AudioSource(48000, 1) track = rtc.LocalAudioTrack.create_audio_track("mic", source) await room.local_participant.publish_track(track) async def _test_wait_disconnect(wait_fn: Callable[[rtc.Room], Awaitable[object]]) -> None: """Connect a room, start waiting, disconnect, and assert RuntimeError.""" name = _room_name() room = rtc.Room() token = _make_token("observer", name) await room.connect(LK_URL, token) task = asyncio.ensure_future(wait_fn(room)) # ensure the task is running and has registered its event handlers await asyncio.sleep(0.5) await room.disconnect() with pytest.raises(RuntimeError, match="disconnected"): await asyncio.wait_for(task, timeout=TIMEOUT) # -- wait_for_participant tests -- @pytest.mark.usefixtures("livekit_server") class TestWaitForParticipant: async def test_joins_after(self): """Room A waits, Room B joins afterwards -> returns B.""" name = _room_name() async with connect_room("observer", name) as room_a: task = asyncio.ensure_future(wait_for_participant(room_a)) await asyncio.sleep(0.5) async with connect_room("joiner", name): result = await asyncio.wait_for(task, timeout=TIMEOUT) assert result.identity == "joiner" async def test_already_joined(self): """Room B already present when A calls wait -> returns immediately.""" name = _room_name() async with connect_room("first", name): await asyncio.sleep(0.5) async with connect_room("observer", name) as room_a: await asyncio.sleep(0.5) result = await asyncio.wait_for(wait_for_participant(room_a), timeout=TIMEOUT) assert result.identity == "first" async def test_with_identity(self): """Only returns participant matching specific identity.""" name = _room_name() async with connect_room("observer", name) as room_a: task = asyncio.ensure_future(wait_for_participant(room_a, identity="target")) await asyncio.sleep(0.3) # Join with wrong identity first async with connect_room("decoy", name): await asyncio.sleep(0.3) assert not task.done() # Now join with the target identity async with connect_room("target", name): result = await asyncio.wait_for(task, timeout=TIMEOUT) assert result.identity == "target" async def test_with_kind(self): """Standard participant doesn't match; agent participant does.""" name = _room_name() async with connect_room("observer", name) as room_a: task = asyncio.ensure_future( wait_for_participant( room_a, kind=rtc.ParticipantKind.PARTICIPANT_KIND_AGENT, ) ) await asyncio.sleep(0.3) # Standard participant should not match async with connect_room("standard", name): await asyncio.sleep(0.3) assert not task.done() # Agent participant should match async with connect_room("agent-p", name, kind="agent", agent=True): result = await asyncio.wait_for(task, timeout=TIMEOUT) assert result.identity == "agent-p" async def test_disconnect(self): """Room disconnects while waiting -> RuntimeError raised.""" await _test_wait_disconnect(wait_for_participant) # -- wait_for_track_publication tests -- @pytest.mark.usefixtures("livekit_server") class TestWaitForTrackPublication: async def test_published_after(self): """Observer waits, publisher publishes audio track -> returns publication.""" name = _room_name() async with connect_room("observer", name) as room_a: task = asyncio.ensure_future(wait_for_track_publication(room_a)) await asyncio.sleep(0.5) async with connect_room("publisher", name) as room_b: await _publish_audio_track(room_b) result = await asyncio.wait_for(task, timeout=TIMEOUT) assert result.kind == rtc.TrackKind.KIND_AUDIO async def test_already_published(self): """Track already published before wait -> returns immediately.""" name = _room_name() async with connect_room("publisher", name) as room_b: await _publish_audio_track(room_b) await asyncio.sleep(0.5) async with connect_room("observer", name) as room_a: await asyncio.sleep(0.5) result = await asyncio.wait_for(wait_for_track_publication(room_a), timeout=TIMEOUT) assert result.kind == rtc.TrackKind.KIND_AUDIO async def test_disconnect(self): """Room disconnects while waiting -> RuntimeError raised.""" await _test_wait_disconnect(wait_for_track_publication) # -- wait_for_agent tests -- @pytest.mark.usefixtures("livekit_server") class TestWaitForAgent: async def test_agent_joins(self): """Agent-kind participant detected correctly.""" name = _room_name() async with connect_room("observer", name) as room_a: task = asyncio.ensure_future(wait_for_agent(room_a)) await asyncio.sleep(0.5) async with connect_room("my-agent", name, kind="agent", agent=True): result = await asyncio.wait_for(task, timeout=TIMEOUT) assert result.identity == "my-agent" assert result.kind == rtc.ParticipantKind.PARTICIPANT_KIND_AGENT # -- Reconnect E2E tests -- # # exercise resume vs full-reconnect behavior end-to-end against a # real LiveKit server. _AGENT_IDENTITY = "agent" _USER_IDENTITY = "user" _TONE_NAME = "agent_tone" # Threshold tuned to discriminate the steady sine wave from server-side # silence/noise. The 440Hz tone at 0.5 amplitude lands well above 0.1. _AUDIO_RMS_THRESHOLD = 0.05 @contextlib.asynccontextmanager async def _agent_publishing_tone(agent_room: rtc.Room): """Publish a steady sine tone via the production `_ParticipantAudioOutput` helper, so the test exercises both the SDK and the agents-framework layer.""" tone = SineToneSource(frequency=440.0, amplitude=0.5) output = _ParticipantAudioOutput( agent_room, sample_rate=tone._sample_rate, num_channels=1, track_publish_options=rtc.TrackPublishOptions(source=rtc.TrackSource.SOURCE_MICROPHONE), track_name=_TONE_NAME, ) # Replace the helper's internal AudioSource with our SineToneSource so the # published track carries the known signal. _ParticipantAudioOutput # constructs its own AudioSource in __init__; swap it before .start() so # the LocalAudioTrack is built against ours. await output._audio_source.aclose() output._audio_source = tone.source await output.start() await tone.start() try: yield output, tone finally: await tone.aclose() await output.aclose() def _agent_audio_publications(user_room: rtc.Room) -> list[rtc.RemoteTrackPublication]: """All audio publications the user sees from the agent.""" agent = user_room.remote_participants.get(_AGENT_IDENTITY) if agent is None: return [] return [ pub for pub in agent.track_publications.values() if pub.kind == rtc.TrackKind.KIND_AUDIO ] async def _await_subscribed_audio(user_room: rtc.Room) -> rtc.RemoteAudioTrack: pub = await asyncio.wait_for( wait_for_track_publication( user_room, identity=_AGENT_IDENTITY, kind=rtc.TrackKind.KIND_AUDIO, wait_for_subscription=True, ), timeout=15.0, ) track = pub.track assert isinstance(track, rtc.RemoteAudioTrack) return track async def _wait_back_to_connected(room: rtc.Room, *, timeout: float = 15.0) -> None: await wait_for_event( room, "connection_state_changed", timeout=timeout, predicate=lambda state: state == rtc.ConnectionState.CONN_CONNECTED, ) @pytest.mark.skipif( not os.environ.get("LIVEKIT_URL"), reason="LIVEKIT_URL not set; skipping reconnect E2E tests " "(set LIVEKIT_URL/LIVEKIT_API_KEY/LIVEKIT_API_SECRET to enable)", ) @pytest.mark.skipif( not hasattr(rtc, "SimulateScenarioKind"), reason="livekit-rtc lacks SimulateScenarioKind; run `make link-rtc-local`", ) class TestReconnect: async def test_resume_preserves_publication_and_audio(self): """Resume must NOT churn the publication set, should fire `reconnected`, and audio must keep flowing.""" room_name = _e2e_room_name("resume") async with ( _connect_e2e_room(_USER_IDENTITY, room_name) as user_room, _connect_e2e_room(_AGENT_IDENTITY, room_name, agent=True) as agent_room, _agent_publishing_tone(agent_room) as (_output, _tone), ): await asyncio.wait_for( wait_for_participant(user_room, identity=_AGENT_IDENTITY), timeout=10.0, ) track = await _await_subscribed_audio(user_room) async with AudioEnergyMonitor.watch(track) as mon: await mon.wait_for_audio(min_rms=_AUDIO_RMS_THRESHOLD, timeout=5.0) # Snapshot publication state before the disturbance. publications_before = _agent_audio_publications(user_room) assert len(publications_before) == 1, publications_before sid_before = publications_before[0].sid reconnected_fired = asyncio.Event() agent_room.on("reconnected", lambda: reconnected_fired.set()) await simulate_resume(agent_room) # Engine cycles through Reconnecting -> Connected; wait for # the second transition. await _wait_back_to_connected(agent_room, timeout=15.0) # Brief grace window for any Reconnected dispatch. await asyncio.sleep(1.0) assert reconnected_fired.is_set(), ( "RoomEvent::Reconnected should be fired on a resume" ) # Publication identity is preserved. publications_after = _agent_audio_publications(user_room) assert len(publications_after) == 1, publications_after assert publications_after[0].sid == sid_before, ( "publication SID changed on resume — engine should have preserved it" ) # Audio continues to flow uninterrupted. await mon.assert_audio_continuous(min_rms=_AUDIO_RMS_THRESHOLD, duration=1.0) async def test_full_reconnect_republishes_once_and_audio_recovers(self): """Full reconnect must fire `reconnected` exactly once, end with exactly one audio publication (the SDK's auto-republish — the agents framework must not produce a duplicate), and audio must recover.""" room_name = _e2e_room_name("full") async with ( _connect_e2e_room(_USER_IDENTITY, room_name) as user_room, _connect_e2e_room(_AGENT_IDENTITY, room_name, agent=True) as agent_room, _agent_publishing_tone(agent_room) as (_output, _tone), ): await asyncio.wait_for( wait_for_participant(user_room, identity=_AGENT_IDENTITY), timeout=10.0, ) track = await _await_subscribed_audio(user_room) async with AudioEnergyMonitor.watch(track) as mon: await mon.wait_for_audio(min_rms=_AUDIO_RMS_THRESHOLD, timeout=5.0) publications_before = _agent_audio_publications(user_room) assert len(publications_before) == 1 # Bug regression check: should fire exactly once. reconnect_count = 0 def _count(*_args) -> None: nonlocal reconnect_count reconnect_count += 1 agent_room.on("reconnected", _count) await simulate_full_reconnect(agent_room) await wait_for_event(agent_room, "reconnected", timeout=20.0) # After the SDK's auto-republish the user observes # unpublish -> publish on the agent. Resubscribe. new_track = await _await_subscribed_audio(user_room) # Bug regression: must be exactly ONE audio publication. # Pre-fix, the agents framework's `_on_reconnected` raced # to publish a second track on top of the SDK's # auto-republished one. await asyncio.sleep(0.5) # let any stray duplicate settle publications_after = _agent_audio_publications(user_room) assert len(publications_after) == 1, ( f"expected exactly 1 audio publication after full reconnect, " f"saw {len(publications_after)}: {[p.sid for p in publications_after]}" ) assert reconnect_count == 1, ( f"reconnected fired {reconnect_count} times; expected exactly 1" ) async with AudioEnergyMonitor.watch(new_track) as new_mon: await new_mon.wait_for_audio(min_rms=_AUDIO_RMS_THRESHOLD, timeout=10.0) await new_mon.assert_audio_continuous( min_rms=_AUDIO_RMS_THRESHOLD, duration=1.0 )