wehub-resource-sync
121 lines
3.8 KiB
Python
121 lines
3.8 KiB
Python
"""Unit tests for xAI TTS websocket behavior."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
from dataclasses import replace
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import aiohttp
|
|
import pytest
|
|
|
|
from livekit.agents import APIConnectOptions
|
|
from livekit.plugins.xai import tts as xai_tts
|
|
|
|
pytestmark = pytest.mark.plugin("xai")
|
|
|
|
|
|
class _FakeWebSocket:
|
|
def __init__(self, audio: bytes = b"audio") -> None:
|
|
self._audio = audio
|
|
self._messages: asyncio.Queue[Any] = asyncio.Queue()
|
|
self.close_code: int | None = None
|
|
self.sent: list[dict[str, Any]] = []
|
|
self.closed = False
|
|
|
|
async def send_str(self, data: str) -> None:
|
|
packet = json.loads(data)
|
|
self.sent.append(packet)
|
|
if packet["type"] == "text.done":
|
|
await self._messages.put(
|
|
SimpleNamespace(
|
|
type=aiohttp.WSMsgType.TEXT,
|
|
data=json.dumps(
|
|
{
|
|
"type": "audio.delta",
|
|
"delta": base64.b64encode(self._audio).decode("ascii"),
|
|
}
|
|
),
|
|
)
|
|
)
|
|
await self._messages.put(
|
|
SimpleNamespace(
|
|
type=aiohttp.WSMsgType.TEXT,
|
|
data=json.dumps({"type": "audio.done"}),
|
|
)
|
|
)
|
|
|
|
async def receive(self) -> Any:
|
|
return await self._messages.get()
|
|
|
|
async def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class _FakeEmitter:
|
|
def __init__(self) -> None:
|
|
self.started_segments: list[str] = []
|
|
self.ended_segments = 0
|
|
self.audio_chunks: list[bytes] = []
|
|
|
|
def start_segment(self, *, segment_id: str) -> None:
|
|
self.started_segments.append(segment_id)
|
|
|
|
def push(self, audio: bytes) -> None:
|
|
self.audio_chunks.append(audio)
|
|
|
|
def end_segment(self) -> None:
|
|
self.ended_segments += 1
|
|
|
|
|
|
def _new_word_stream(tts: xai_tts.TTS, text: str):
|
|
word_stream = tts._opts.tokenizer.stream() # pyright: ignore[reportPrivateUsage]
|
|
word_stream.push_text(text)
|
|
word_stream.end_input()
|
|
return word_stream
|
|
|
|
|
|
def _new_synthesize_stream(tts: xai_tts.TTS) -> xai_tts.SynthesizeStream:
|
|
stream = object.__new__(xai_tts.SynthesizeStream)
|
|
stream._tts = tts # pyright: ignore[reportPrivateUsage]
|
|
stream._opts = replace(tts._opts) # pyright: ignore[reportPrivateUsage]
|
|
stream._conn_options = APIConnectOptions(max_retry=0, timeout=1.0) # pyright: ignore[reportPrivateUsage]
|
|
stream._started_time = 0.0 # pyright: ignore[reportPrivateUsage]
|
|
stream._acquire_time = 0.0 # pyright: ignore[reportPrivateUsage]
|
|
stream._connection_reused = False # pyright: ignore[reportPrivateUsage]
|
|
return stream
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_segments_reuse_websocket_connection(monkeypatch: pytest.MonkeyPatch):
|
|
tts = xai_tts.TTS(api_key="test-key")
|
|
stream = _new_synthesize_stream(tts)
|
|
emitter = _FakeEmitter()
|
|
websockets: list[_FakeWebSocket] = []
|
|
|
|
async def connect_ws(*_args: object, **_kwargs: object) -> _FakeWebSocket:
|
|
ws = _FakeWebSocket()
|
|
websockets.append(ws)
|
|
return ws
|
|
|
|
async def close_ws(ws: _FakeWebSocket) -> None:
|
|
await ws.close()
|
|
|
|
monkeypatch.setattr(tts, "_connect_ws", connect_ws)
|
|
monkeypatch.setattr(tts, "_close_ws", close_ws)
|
|
|
|
await stream._run_ws(_new_word_stream(tts, "first"), emitter) # pyright: ignore[reportPrivateUsage]
|
|
await stream._run_ws(_new_word_stream(tts, "second"), emitter) # pyright: ignore[reportPrivateUsage]
|
|
|
|
assert len(websockets) == 1
|
|
assert websockets[0].closed is False
|
|
assert emitter.audio_chunks == [b"audio", b"audio"]
|
|
assert emitter.ended_segments == 2
|
|
|
|
await tts.aclose()
|
|
|
|
assert websockets[0].closed is True
|