first vommit

This commit is contained in:
Till-Immer committed 2026-07-04 14:31:50 +02:00
commit a04e68bc95
4926 files changed
+764440

No files matched your search

@@ -0,0 +1,16 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,389 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Tests for live_music.py."""
import contextlib
import json
from typing import AsyncIterator
from unittest import mock
from unittest.mock import AsyncMock
from unittest.mock import Mock
from unittest.mock import patch
import warnings
from google.oauth2.credentials import Credentials
import pytest
from websockets import client
from ... import _api_client as api_client
from ... import _common
from ... import Client
from ... import client as gl_client
from ... import live
from ... import live_music
from ... import types
from .. import pytest_helper
try:
import aiohttp
AIOHTTP_NOT_INSTALLED = False
except ImportError:
AIOHTTP_NOT_INSTALLED = True
aiohttp = mock.MagicMock()
requires_aiohttp = pytest.mark.skipif(
AIOHTTP_NOT_INSTALLED, reason="aiohttp is not installed, skipping test."
)
def mock_api_client(vertexai=False, credentials=None):
api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
if not vertexai:
api_client.api_key = 'TEST_API_KEY'
api_client.location = None
api_client.project = None
else:
api_client.api_key = None
api_client.location = 'us-central1'
api_client.project = 'test_project'
api_client._host = lambda: 'test_host'
api_client._credentials = credentials
api_client._http_options = types.HttpOptions.model_validate(
{'headers': {}}
) # Ensure headers exist
api_client.vertexai = vertexai
api_client._api_client = api_client
return api_client
@pytest.fixture
def mock_websocket():
websocket = AsyncMock(spec=client.ClientConnection)
websocket.send = AsyncMock()
websocket.recv = AsyncMock(
return_value=b"""{
"serverContent": {
"audioChunks": [
{
"data": "Z2VsYmFuYW5h",
"mimeType": "audio/l16;rate=48000;channels=2",
"sourceMetadata": {
"clientContent": {
"weightedPrompts": [
{
"text": "Jazz",
"weight": 1
}
]
},
"musicGenerationConfig": {
"seed": -957124937,
"bpm": 140,
"scale": "A_FLAT_MAJOR_F_MINOR"
}
}
}
]
}
}"""
) # Default response
websocket.close = AsyncMock()
return websocket
async def get_connect_message(api_client, model):
mock_ws = AsyncMock()
mock_ws.send = AsyncMock()
mock_ws.recv = AsyncMock(return_value=b'some response')
mock_google_auth_default = Mock(return_value=(None, None))
mock_creds = Mock(token='test_token')
mock_google_auth_default.return_value = (mock_creds, None)
@contextlib.asynccontextmanager
async def mock_connect(uri, additional_headers=None):
yield mock_ws
@patch('google.auth.default', new=mock_google_auth_default)
@patch.object(live_music, 'connect', new=mock_connect)
async def _test_connect():
live_module = live.AsyncLive(api_client)
async with live_module.music.connect(
model=model,
):
pass
mock_ws.send.assert_called_once()
return json.loads(mock_ws.send.call_args[0][0])
return await _test_connect()
@pytest.mark.asyncio
async def test_connect_uses_header_auth_without_query_key(mock_websocket):
api_client = mock_api_client(vertexai=False)
api_client._websocket_base_url = lambda: 'wss://generativelanguage.googleapis.com'
api_client._http_options.api_version = 'v1beta'
api_client._http_options.headers['x-goog-api-key'] = 'TEST_API_KEY'
captured = {}
@contextlib.asynccontextmanager
async def mock_connect(uri, additional_headers=None):
captured['uri'] = uri
captured['headers'] = additional_headers
yield mock_websocket
@patch.object(live_music, 'connect', new=mock_connect)
async def _test_connect():
live_module = live.AsyncLive(api_client)
async with live_module.music.connect(model='test_model'):
pass
await _test_connect()
assert 'TEST_API_KEY' not in captured['uri']
assert '?key=' not in captured['uri']
assert captured['headers']['x-goog-api-key'] == 'TEST_API_KEY'
def test_mldev_from_env(monkeypatch):
api_key = 'google_api_key'
monkeypatch.setenv('GOOGLE_API_KEY', api_key)
client = Client()
assert not client.aio.live.music._api_client.vertexai
assert client.aio.live.music._api_client.api_key == api_key
assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
@requires_aiohttp
def test_vertex_from_env(monkeypatch):
project_id = 'fake_project_id'
location = 'fake-location'
monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', 'true')
monkeypatch.setenv('GOOGLE_CLOUD_PROJECT', project_id)
monkeypatch.setenv('GOOGLE_CLOUD_LOCATION', location)
client = Client()
assert client.aio.live.music._api_client.vertexai
assert client.aio.live.music._api_client.project == project_id
assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
def test_websocket_base_url():
base_url = 'https://test.com'
api_client = gl_client.BaseApiClient(
api_key='google_api_key',
http_options={'base_url': base_url},
)
assert api_client._websocket_base_url() == 'wss://test.com'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_weighted_prompts(
mock_websocket, vertexai
):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
await session.set_weighted_prompts(prompts=[types.WeightedPrompt(text='Jazz', weight=1)])
return
await session.set_weighted_prompts(prompts=[types.WeightedPrompt(text='Jazz', weight=1)])
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'clientContent' in sent_data
assert sent_data['clientContent']['weightedPrompts'][0]['text'] == 'Jazz'
assert sent_data['clientContent']['weightedPrompts'][0]['weight'] == 1
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_config(
mock_websocket, vertexai
):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
await session.set_music_generation_config(
config=types.LiveMusicGenerationConfig(
bpm=140,
music_generation_mode=types.MusicGenerationMode.VOCALIZATION,
)
)
return
await session.set_music_generation_config(
config=types.LiveMusicGenerationConfig(
bpm=140,
music_generation_mode=types.MusicGenerationMode.VOCALIZATION,
)
)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'musicGenerationConfig' in sent_data
assert sent_data['musicGenerationConfig']['bpm'] == 140
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_control_signal_play(
mock_websocket, vertexai
):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
await session.play()
return
await session.play()
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'playbackControl' in sent_data
assert 'PLAY' in sent_data['playbackControl']
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_control_signal_pause(
mock_websocket, vertexai
):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
await session.pause()
return
await session.pause()
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'playbackControl' in sent_data
assert 'PAUSE' in sent_data['playbackControl']
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_control_signal_stop(
mock_websocket, vertexai
):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
await session.stop()
return
await session.stop()
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'playbackControl' in sent_data
assert 'STOP' in sent_data['playbackControl']
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_control_signal_reset_context(
mock_websocket, vertexai
):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
await session.reset_context()
return
await session.reset_context()
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'playbackControl' in sent_data
assert 'RESET_CONTEXT' in sent_data['playbackControl']
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_receive( mock_websocket, vertexai):
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
if vertexai:
with pytest.raises(NotImplementedError):
async for _ in session.receive():
pass
return
async for response in session.receive():
assert isinstance(response, types.LiveMusicServerMessage)
audio_chunk = response.server_content.audio_chunks[0]
# Data contains decoded b64 audio
assert audio_chunk.data == b'gelbanana'
assert audio_chunk.mime_type == 'audio/l16;rate=48000;channels=2'
assert audio_chunk.source_metadata.client_content.weighted_prompts[0].text == 'Jazz'
assert audio_chunk.source_metadata.client_content.weighted_prompts[0].weight == 1
assert audio_chunk.source_metadata.music_generation_config.bpm == 140
break
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_receive_error(
mock_websocket, vertexai
):
mock_websocket.recv = AsyncMock(return_value='invalid json')
session = live_music.AsyncMusicSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
with pytest.raises(ValueError):
await session.receive().__anext__()
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_close( mock_websocket, vertexai):
session = live_music.AsyncMusicSession(
mock_api_client(vertexai=vertexai), mock_websocket
)
await session.close()
mock_websocket.close.assert_called_once()
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_setup_to_api(vertexai):
if vertexai:
with pytest.raises(NotImplementedError):
await get_connect_message(
mock_api_client(vertexai=vertexai),
model='test_model'
)
return
result = await get_connect_message(
mock_api_client(vertexai=vertexai),
model='test_model'
)
expected_result = {'setup': {}}
if vertexai:
# Vertex is not supported yet
assert False
else:
expected_result['setup']['model'] = 'models/test_model'
assert result == expected_result
@@ -0,0 +1,163 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Tests for live response handling."""
import json
from typing import cast
from unittest import mock
from unittest.mock import AsyncMock
import pytest
from ... import _api_client as api_client
from ... import _common
from ... import Client
from ... import client as gl_client
from ... import live
from ... import types
def mock_api_client(vertexai=False):
"""Creates a mock BaseApiClient."""
mock_client = mock.MagicMock(spec=gl_client.BaseApiClient)
if not vertexai:
mock_client.api_key = 'TEST_API_KEY'
mock_client.location = None
mock_client.project = None
else:
mock_client.api_key = None
mock_client.location = 'us-central1'
mock_client.project = 'test_project'
mock_client._host = lambda: 'test_host'
mock_client._http_options = types.HttpOptions.model_validate(
{'headers': {}}
)
mock_client.vertexai = vertexai
return mock_client
@pytest.fixture
def mock_websocket():
"""Provides a mock websocket connection."""
# Use live.ClientConnection if that's the specific type hint in AsyncSession
websocket = AsyncMock(spec=live.ClientConnection)
websocket.send = AsyncMock()
# Set default recv value, will be overridden in the test
websocket.recv = AsyncMock(return_value='{}')
websocket.close = AsyncMock()
return websocket
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_receive_server_content(mock_websocket, vertexai):
raw_response_json = json.dumps({
"usageMetadata": {
"promptTokenCount": 15,
"responseTokenCount": 25,
"candidatesTokenCount": 50,
"totalTokenCount": 200,
"responseTokensDetails": [
{
"tokenCount": 20,
"modality": "TEXT",
}
],
"candidatesTokensDetails": [
{
"tokenCount": 10,
"modality": "TEXT",
}
],
},
"serverContent": {
"modelTurn": {
"parts": [{"text": "This is a simple response."}]
},
"turnComplete": True,
"groundingMetadata": {
"web_search_queries": ["test query"],
"groundingChunks": [{
"web": {
"domain": "google.com",
"title": "Search results",
}
}]
}
}
})
mock_websocket.recv.return_value = raw_response_json
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
result = await session._receive()
# Assert the results
assert isinstance(result, types.LiveServerMessage)
assert (
result.server_content.model_turn.parts[0].text
== "This is a simple response."
)
assert result.server_content.turn_complete
assert result.server_content.grounding_metadata.web_search_queries == ["test query"]
assert result.server_content.grounding_metadata.grounding_chunks[0].web.domain == "google.com"
assert result.server_content.grounding_metadata.grounding_chunks[0].web.title == "Search results"
# Verify usageMetadata was parsed
assert isinstance(result.usage_metadata, types.UsageMetadata)
assert result.usage_metadata.prompt_token_count == 15
assert result.usage_metadata.total_token_count == 200
if not vertexai:
assert result.usage_metadata.response_token_count == 25
assert result.usage_metadata.response_tokens_details[0].token_count == 20
else:
# VertexAI maps candidatesTokenCount to responseTokenCount and maps
# candidatesTokensDetails to responseTokensDetails.
assert result.usage_metadata.response_token_count == 50
assert result.usage_metadata.response_tokens_details[0].token_count == 10
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_receive_server_content_with_turn_reason(mock_websocket, vertexai):
"""Tests parsing of LiveServerContent with turn_complete_reason and waiting_for_input."""
raw_response_json = json.dumps({
"serverContent": {
"modelTurn": {
"parts": [{"text": "Please provide more details."}]
},
"turnComplete": True,
"turnCompleteReason": "NEED_MORE_INPUT",
"waitingForInput": True
}
})
mock_websocket.recv.return_value = raw_response_json
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
result = await session._receive()
# Assert the results
assert isinstance(result, types.LiveServerMessage)
assert result.server_content is not None
assert result.server_content.model_turn.parts[0].text == "Please provide more details."
assert result.server_content.turn_complete is True
assert result.server_content.turn_complete_reason == types.TurnCompleteReason.NEED_MORE_INPUT
assert result.server_content.waiting_for_input is True
@@ -0,0 +1,147 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Tests for live.py."""
import base64
import json
from unittest import mock
import pytest
from websockets import client
from .. import pytest_helper
from ... import client as gl_client
from ... import live
from ... import types
def mock_api_client(vertexai=False):
api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
api_client.api_key = 'TEST_API_KEY'
api_client._host = lambda: 'test_host'
api_client._http_options = {'headers': {}} # Ensure headers exist
api_client.vertexai = vertexai
return api_client
@pytest.fixture
def mock_websocket():
websocket = mock.AsyncMock(spec=client.ClientConnection)
websocket.send = mock.AsyncMock()
websocket.recv = mock.AsyncMock(
return_value='{"serverContent": {"turnComplete": true}}'
) # Default response
websocket.close = mock.AsyncMock()
return websocket
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_content_dict(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
content = [{'parts': [{'text': 'test'}]}]
await session.send_client_content(turns=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'client_content' in sent_data
assert sent_data['client_content']['turns'][0]['parts'][0]['text'] == 'test'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_content_dict_list(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
content = [{'parts': [{'text': 'test'}]}]
await session.send_client_content(turns=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'client_content' in sent_data
assert sent_data['client_content']['turns'][0]['parts'][0]['text'] == 'test'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_content_content(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
content = types.Content.model_validate({'parts': [{'text': 'test'}]})
await session.send_client_content(turns=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'client_content' in sent_data
assert sent_data['client_content']['turns'][0]['parts'][0]['text'] == 'test'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_content_with_blob(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
content = types.Content.model_validate(
{'parts': [{'inline_data': {'data': b'test'}}]}
)
await session.send_client_content(turns=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'client_content' in sent_data
assert pytest_helper.get_value_ignore_key_case(
sent_data['client_content']['turns'][0]['parts'][0], 'inline_data') == {
'data': base64.b64encode(b'test').decode()
}
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_client_content_turn_complete_false(
mock_websocket, vertexai
):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
await session.send_client_content(turn_complete=False)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'client_content' in sent_data
assert sent_data['client_content']['turnComplete'] == False
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_client_content_empty(
mock_websocket, vertexai
):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
await session.send_client_content()
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'client_content' in sent_data
assert sent_data['client_content']['turnComplete'] == True
@@ -0,0 +1,268 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Tests for live.py."""
import json
import os
from unittest import mock
import pytest
from websockets import client
from ... import client as gl_client
from ... import live
from ... import types
from .. import pytest_helper
IMAGE_FILE_PATH = os.path.abspath(
os.path.join(os.path.dirname(__file__), '../data/google.jpg')
)
def mock_api_client(vertexai=False):
api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
api_client.api_key = 'TEST_API_KEY'
api_client._host = lambda: 'test_host'
api_client._http_options = {'headers': {}} # Ensure headers exist
api_client.vertexai = vertexai
api_client._api_client = api_client
return api_client
@pytest.fixture
def mock_websocket():
websocket = mock.AsyncMock(spec=client.ClientConnection)
websocket.send = mock.AsyncMock()
websocket.recv = mock.AsyncMock(
return_value='{"serverContent": {"turnComplete": true}}'
) # Default response
websocket.close = mock.AsyncMock()
return websocket
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_media_blob_dict(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
content = {'data': bytes([0, 0, 0, 0, 0, 0]), 'mime_type': 'audio/pcm'}
await session.send_realtime_input(media=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['mediaChunks'][0]['data'] == 'AAAAAAAA'
assert pytest_helper.get_value_ignore_key_case(
sent_data['realtime_input']['mediaChunks'][0], 'mime_type'
) == 'audio/pcm'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_media_blob(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
content = types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='audio/pcm')
await session.send_realtime_input(media=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['mediaChunks'][0]['data'] == 'AAAAAAAA'
assert pytest_helper.get_value_ignore_key_case(
sent_data['realtime_input']['mediaChunks'][0], 'mime_type'
) == 'audio/pcm'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_media_image(mock_websocket, vertexai, image_jpeg):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
await session.send_realtime_input(media=image_jpeg)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert pytest_helper.get_value_ignore_key_case(
sent_data['realtime_input']['mediaChunks'][0], 'mime_type'
) == 'image/jpeg'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.parametrize(
'content',
[
{'data': bytes([0, 0, 0, 0, 0, 0]), 'mime_type': 'audio/pcm'},
types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='audio/pcm'),
],
)
@pytest.mark.asyncio
async def test_send_audio(mock_websocket, vertexai, content):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
await session.send_realtime_input(audio=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['audio']['data'] == 'AAAAAAAA'
assert pytest_helper.get_value_ignore_key_case(
sent_data['realtime_input']['audio'], 'mime_type'
) == 'audio/pcm'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_bad_audio_blob(mock_websocket, vertexai):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
content = types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='image/png')
with pytest.raises(ValueError, match='.*Unsupported mime type.*'):
await session.send_realtime_input(audio=content)
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_bad_video_blob(mock_websocket, vertexai):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
content = types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='audio/pcm')
with pytest.raises(ValueError, match='.*Unsupported mime type.*'):
await session.send_realtime_input(video=content)
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_audio_stream_end(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
await session.send_realtime_input(audio_stream_end=True)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['audioStreamEnd'] == True
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.parametrize(
'content',
[
{'data': bytes([0, 0, 0, 0, 0, 0]), 'mime_type': 'image/png'},
types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='image/png'),
],
)
@pytest.mark.asyncio
async def test_send_video(mock_websocket, vertexai, content):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
await session.send_realtime_input(video=content)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['video']['data'] == 'AAAAAAAA'
assert pytest_helper.get_value_ignore_key_case(
sent_data['realtime_input']['video'], 'mime_type'
) == 'image/png'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_video_image(mock_websocket, vertexai, image_jpeg):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
await session.send_realtime_input(video=image_jpeg)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert pytest_helper.get_value_ignore_key_case(
sent_data['realtime_input']['video'], 'mime_type'
) == 'image/jpeg'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_text(mock_websocket, vertexai):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
await session.send_realtime_input(text='Hello?')
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['text'] == 'Hello?'
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.parametrize('activity', [{}, types.ActivityStart()])
@pytest.mark.asyncio
async def test_send_activity_start(mock_websocket, vertexai, activity):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
await session.send_realtime_input(activity_start=activity)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['activityStart'] == {}
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.parametrize('activity', [{}, types.ActivityEnd()])
@pytest.mark.asyncio
async def test_send_activity_end(mock_websocket, vertexai, activity):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
await session.send_realtime_input(activity_end=activity)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'realtime_input' in sent_data
assert sent_data['realtime_input']['activityEnd'] == {}
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_send_multiple_args(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
with pytest.raises(ValueError, match='.*one argument.*'):
await session.send_realtime_input(
text='Hello?', activity_start=types.ActivityStart()
)
@@ -0,0 +1,222 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Tests for live.py."""
import contextlib
import json
from unittest import mock
import pytest
from websockets import client
from .. import pytest_helper
from ... import client as gl_client
from ... import live
from ... import types
def exception_if_mldev(vertexai, exception_type: type[Exception]):
if vertexai:
return contextlib.nullcontext()
else:
return pytest.raises(exception_type)
def mock_api_client(vertexai=False):
api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
api_client.api_key = 'TEST_API_KEY'
api_client._host = lambda: 'test_host'
api_client._http_options = {'headers': {}} # Ensure headers exist
api_client.vertexai = vertexai
api_client._api_client = api_client
return api_client
@pytest.fixture
def mock_websocket():
websocket = mock.AsyncMock(spec=client.ClientConnection)
websocket.send = mock.AsyncMock()
websocket.recv = mock.AsyncMock(
return_value='{"serverContent": {"turnComplete": true}}'
) # Default response
websocket.close = mock.AsyncMock()
return websocket
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_function_response_dict(mock_websocket, vertexai):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(
api_client=api_client, websocket=mock_websocket
)
input = {
'name': 'get_current_weather',
'response': {'temperature': 14.5, 'unit': 'C'},
}
if not vertexai:
input['id'] = 'some-id'
await session.send_tool_response(function_responses=input)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'tool_response' in sent_data
assert (
sent_data['tool_response']['functionResponses'][0]['name']
== 'get_current_weather'
)
assert (
sent_data['tool_response']['functionResponses'][0]['response'][
'temperature'
]
== 14.5
)
assert (
sent_data['tool_response']['functionResponses'][0]['response']['unit']
== 'C'
)
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_function_response(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
input = types.FunctionResponse(
name='get_current_weather',
response={
'temperature': 14.5,
'unit': 'C',
'user_name': 'test_user_name',
'userEmail': 'test_user_email',
},
)
if not vertexai:
input.id = 'some-id'
await session.send_tool_response(function_responses=input)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'tool_response' in sent_data
assert (
sent_data['tool_response']['functionResponses'][0]['name']
== 'get_current_weather'
)
assert (
sent_data['tool_response']['functionResponses'][0]['response']
== input.response
)
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_function_response_scheduling(mock_websocket, vertexai):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
input = types.FunctionResponse(
name='get_current_weather',
response={'temperature': 14.5, 'unit': 'C'},
will_continue=True,
scheduling=types.FunctionResponseScheduling.SILENT,
)
if not vertexai:
input.id = 'some-id'
await session.send_tool_response(function_responses=input)
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'tool_response' in sent_data
assert pytest_helper.get_value_ignore_key_case(
sent_data['tool_response']['functionResponses'][0], 'will_continue'
)
assert (
sent_data['tool_response']['functionResponses'][0]['scheduling']
== 'SILENT'
)
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_function_response_list(mock_websocket, vertexai):
session = live.AsyncSession(
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
)
input1 = {
'name': 'get_current_weather',
'response': {'temperature': 14.5, 'unit': 'C'},
}
input2 = {
'name': 'get_current_weather',
'response': {'temperature': 99.9, 'unit': 'C'},
}
if not vertexai:
input1['id'] = '1'
input2['id'] = '2'
await session.send_tool_response(function_responses=[input1, input2])
mock_websocket.send.assert_called_once()
sent_data = json.loads(mock_websocket.send.call_args[0][0])
assert 'tool_response' in sent_data
assert len(sent_data['tool_response']['functionResponses']) == 2
assert (
sent_data['tool_response']['functionResponses'][0]['response'][
'temperature'
]
== 14.5
)
assert (
sent_data['tool_response']['functionResponses'][1]['response'][
'temperature'
]
== 99.9
)
@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_missing_id(mock_websocket, vertexai):
api_client = mock_api_client(vertexai=vertexai)
session = live.AsyncSession(
api_client=api_client, websocket=mock_websocket
)
input1 = {
'name': 'get_current_weather',
'response': {'temperature': 14.5, 'unit': 'C'},
'id': '1',
}
input2 = {
'name': 'get_current_weather',
'response': {'temperature': 99.9, 'unit': 'C'},
}
if not vertexai:
with pytest.raises(ValueError, match=".*must have.*"):
await session.send_tool_response(function_responses=[input1, input2])