wehub-resource-sync
65 lines
1.8 KiB
Python
65 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from livekit.agents import tts
|
|
from livekit.plugins.google.beta.gemini_tts import TTS
|
|
|
|
pytestmark = pytest.mark.plugin("google")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("livekit.plugins.google.beta.gemini_tts.Client")
|
|
async def test_gemini_tts_success(mock_genai_client_class) -> None:
|
|
# Setup mocks for GenAI Client
|
|
mock_client = MagicMock()
|
|
mock_genai_client_class.return_value = mock_client
|
|
|
|
mock_stream = AsyncMock()
|
|
mock_client.aio.models.generate_content_stream = mock_stream
|
|
|
|
# Mock chunk response candidates
|
|
class MockInlineData:
|
|
def __init__(self, data: bytes):
|
|
self.data = data
|
|
self.mime_type = "audio/pcm"
|
|
|
|
class MockPart:
|
|
def __init__(self, data: bytes):
|
|
self.inline_data = MockInlineData(data)
|
|
|
|
class MockContent:
|
|
def __init__(self, data: bytes):
|
|
self.parts = [MockPart(data)]
|
|
|
|
class MockCandidate:
|
|
def __init__(self, data: bytes):
|
|
self.content = MockContent(data)
|
|
|
|
class MockChunk:
|
|
def __init__(self, data: bytes):
|
|
self.candidates = [MockCandidate(data)]
|
|
|
|
async def mock_generator(*args, **kwargs):
|
|
yield MockChunk(b"\x00" * 4800)
|
|
yield MockChunk(b"\x01" * 4800)
|
|
|
|
mock_stream.side_effect = mock_generator
|
|
|
|
# Initialize TTS
|
|
google_tts = TTS(api_key="test-api-key")
|
|
|
|
# Create output emitter mock
|
|
mock_emitter = MagicMock(spec=tts.AudioEmitter)
|
|
|
|
# Run ChunkedStream
|
|
stream = google_tts.synthesize("Hello world")
|
|
await stream._run(mock_emitter)
|
|
|
|
# Assertions
|
|
mock_stream.assert_called_once()
|
|
mock_emitter.initialize.assert_called_once()
|
|
assert mock_emitter.push.call_count == 2
|