# 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
import os
import ssl
import typing
from typing import Any, AsyncIterator
from unittest import mock
from unittest.mock import AsyncMock
from unittest.mock import Mock
from unittest.mock import patch
import warnings

import certifi
from google.oauth2.credentials import Credentials
import pytest
from websockets import client

from .. import pytest_helper
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
try:
    import aiohttp
    AIOHTTP_NOT_INSTALLED = False
except ImportError:
    AIOHTTP_NOT_INSTALLED = True
    aiohttp = mock.MagicMock()


if typing.TYPE_CHECKING:
  from mcp import types as mcp_types
  from mcp import ClientSession as McpClientSession
else:
  mcp_types: typing.Type = Any
  McpClientSession: typing.Type = Any
  try:
    from mcp import types as mcp_types
    from mcp import ClientSession as McpClientSession
  except ImportError:
    mcp_types = None
    McpClientSession = None


requires_aiohttp = pytest.mark.skipif(
    AIOHTTP_NOT_INSTALLED, reason="aiohttp is not installed, skipping test."
)

function_declarations = [{
    'name': 'get_current_weather',
    'description': 'Get the current weather in a city',
    'parameters': {
        'type': 'OBJECT',
        'properties': {
            'location': {
                'type': 'STRING',
                'description': 'The location to get the weather for',
            },
            'unit': {
                'type': 'STRING',
                'enum': ['C', 'F'],
            },
        },
    },
}]


def get_current_weather(location: str, unit: str):
  """Get the current weather in a city."""
  return 15 if unit == 'C' else 59


def mock_api_client(vertexai=False, credentials=None, http_options=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
    api_client.custom_base_url = None
  else:
    api_client.api_key = None
    if http_options:
      http_options = (
          types.HttpOptions(**http_options)
          if isinstance(http_options, dict)
          else http_options
      )
      api_client.custom_base_url = http_options.base_url
      api_client.location = None
      api_client.project = None
    else:
      api_client.location = 'us-central1'
      api_client.project = 'test_project'
      api_client.custom_base_url = None

  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
  ctx = ssl.create_default_context(
      cafile=os.environ.get("SSL_CERT_FILE", certifi.where()),
      capath=os.environ.get("SSL_CERT_DIR"),
  )
  api_client._websocket_ssl_ctx = {'ssl': ctx}
  return api_client


@pytest.fixture
def mock_websocket():
  websocket = AsyncMock(spec=client.ClientConnection)
  websocket.send = AsyncMock()
  websocket.recv = AsyncMock(
      return_value='{"serverContent": {"turnComplete": true}}'
  )  # Default response
  websocket.close = AsyncMock()
  return websocket


async def get_connect_message(api_client, model, config=None):
  if config is None:
    config = {}
  mock_ws = AsyncMock()
  mock_ws.send = AsyncMock()
  mock_ws.recv = AsyncMock(
      return_value=(
          b'{\n  "setupComplete": {"sessionId": "test_session_id"}\n}\n'
      )
  )

  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, **kwargs):
    yield mock_ws

  @patch('google.auth.default', new=mock_google_auth_default)
  @patch.object(live, 'ws_connect', new=mock_connect)
  async def _test_connect():
    live_module = live.AsyncLive(api_client)
    async with live_module.connect(
        model=model,
        config=config,
    ):
      pass

    mock_ws.send.assert_called_once()
    return json.loads(mock_ws.send.call_args[0][0])

  return await _test_connect()


async def _async_iterator_to_list(async_iter):
  return [value async for value in async_iter]


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._api_client.vertexai
  assert client.aio.live._api_client.api_key == api_key
  assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
  assert client.aio.live._api_client._http_options.headers['x-goog-api-key'] == api_key


@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._api_client.custom_base_url is None
  assert client.aio.live._api_client.vertexai
  assert client.aio.live._api_client.project == project_id
  assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
  assert 'x-goog-api-key' not in client.aio.live._api_client._http_options.headers


def test_vertex_api_key_from_env(monkeypatch):
  api_key = 'google_api_key'
  monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', 'true')
  monkeypatch.setenv('GOOGLE_API_KEY', api_key)

  # Due to proj/location taking precedence, need to clear proj/location env
  # variables. Tests in client/test_client_initialization.py provide
  # comprehensive coverage for proj/location and api key precedence.
  monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "")
  monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "")

  client = Client()

  assert client.aio.live._api_client.vertexai
  assert client.aio.live._api_client.api_key == api_key
  assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
  assert client.aio.live._api_client._http_options.headers['x-goog-api-key'] == api_key


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'


def test_websocket_base_url_no_auth_with_custom_base_url():
  base_url = 'https://test-api-gateway-proxy.com'
  api_client = gl_client.BaseApiClient(
      vertexai=True,
      http_options={
          'base_url': base_url,
          'headers': {'Authorization': 'Bearer test_token'},
      },
  )
  # Note that our test environment does have project/location set. So we
  # need to explicitly set them to None here.
  api_client.project = None
  api_client.location = None

  # Fully pass the custom base url if no API key or project/location.
  assert api_client._websocket_base_url() == base_url


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_text(mock_websocket, vertexai):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  await session.send(input='test')
  mock_websocket.send.assert_called_once()
  sent_data = json.loads(mock_websocket.send.call_args[0][0])
  assert 'client_content' in sent_data


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_content_dict(
    mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  client_content = {
      'content': [{'parts': [{'text': 'test'}]}],
      'turn_complete': True,
  }
  await session.send(input=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


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_content(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  client_content = types.LiveClientContent(
      turns=[types.Content(parts=[types.Part(text='test')])], turn_complete=True
  )
  await session.send(input=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


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_bytes(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  realtime_input = {'data': b'000000', 'mime_type': 'audio/pcm'}

  await session.send(input=realtime_input)
  mock_websocket.send.assert_called_once()
  sent_data = json.loads(mock_websocket.send.call_args[0][0])
  assert 'realtime_input' in sent_data


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_blob(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  realtime_input = types.Blob(data=b'000000', mime_type='audio/pcm')

  await session.send(input=realtime_input)
  mock_websocket.send.assert_called_once()
  sent_data = json.loads(mock_websocket.send.call_args[0][0])
  assert 'realtime_input' in sent_data


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_realtime_input(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  realtime_input = types.LiveClientRealtimeInput(
      media_chunks=[types.Blob(data='MDAwMDAw', mime_type='audio/pcm')]
  )
  await session.send(input=realtime_input)
  mock_websocket.send.assert_called_once()
  sent_data = json.loads(mock_websocket.send.call_args[0][0])
  assert 'realtime_input' in sent_data


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_tool_response(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )

  if vertexai:
    tool_response = types.LiveClientToolResponse(
        function_responses=[
            types.FunctionResponse(
                name='get_current_weather',
                response={'temperature': 14.5, 'unit': 'C'},
            )
        ]
    )
  else:
    tool_response = types.LiveClientToolResponse(
        function_responses=[
            types.FunctionResponse(
                name='get_current_weather',
                response={'temperature': 14.5, 'unit': 'C'},
                id='some-id',
            )
        ]
    )
  await session.send(input=tool_response)
  mock_websocket.send.assert_called_once()
  sent_data = json.loads(mock_websocket.send.call_args[0][0])
  assert 'tool_response' in sent_data


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_input_none(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  await session.send(input=None)
  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']['turn_complete']


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_send_error(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  with pytest.raises(ValueError):
    await session.send(input=[{'invalid_key': 'invalid_value'}])

  with pytest.raises(ValueError):
    await session.send(input={'invalid_key': 'invalid_value'})


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_receive( mock_websocket, vertexai):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  responses = session.receive()
  responses = await _async_iterator_to_list(responses)
  assert isinstance(responses[0], types.LiveServerMessage)


@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.AsyncSession(
      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_receive_text(
     mock_websocket, vertexai
):
  mock_websocket.recv = AsyncMock(
      side_effect=[
          '{"serverContent": {"modelTurn": {"parts":[{"text": "test"}]}}}',
          '{"serverContent": {"turnComplete": true}}',
      ]
  )
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  messages = session.receive()
  messages = await _async_iterator_to_list(messages)
  assert isinstance(messages[0], types.LiveServerMessage)
  assert messages[0].server_content.model_turn.parts[0].text == 'test'
  assert messages[1].server_content.turn_complete == True


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_receive_audio(
     mock_websocket, vertexai
):
  mock_websocket.recv = AsyncMock(
      side_effect=[
          (
              '{"serverContent": {"modelTurn": {"parts":[{"inlineData":'
              ' {"data": "MDAwMDAw", "mimeType": "audio/pcm" }}]}}}'
          ),
          '{"serverContent": {"turnComplete": true}}',
      ]
  )
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  messages = session.receive()
  messages = await _async_iterator_to_list(messages)
  assert isinstance(messages[0], types.LiveServerMessage)
  assert (
      messages[0].server_content.model_turn.parts[0].inline_data.mime_type
      == 'audio/pcm'
  )
  assert (
      messages[0].server_content.model_turn.parts[0].inline_data.data
      == b'000000'
  )

  with pytest.raises(RuntimeError):
    await _async_iterator_to_list(session.receive())


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_receive_tool_call(
     mock_websocket, vertexai
):
  mock_websocket.recv = AsyncMock(
      side_effect=[
          (
              '{"toolCall": {"functionCalls": [{"name":'
              ' "get_current_weather", "args": {"location": "San Francisco",'
              ' "unit": "C"}}]}}'
          ),
          '{"serverContent": {"turnComplete": true}}',
      ]
  )
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  messages = session.receive()
  messages = await _async_iterator_to_list(messages)
  assert isinstance(messages[0], types.LiveServerMessage)
  assert messages[0].tool_call.function_calls[0].name == 'get_current_weather'
  assert (
      messages[0].tool_call.function_calls[0].args['location']
      == 'San Francisco'
  )
  assert messages[0].tool_call.function_calls[0].args['unit'] == 'C'

  with pytest.raises(RuntimeError):
    await _async_iterator_to_list(session.receive())


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_receive_transcription(
     mock_websocket, vertexai
):
  mock_websocket.recv = AsyncMock(
      side_effect=[
          '{"serverContent": {"inputTranscription": {"text": "test_input", "finished": true}}}',
          '{"serverContent": {"outputTranscription": {"text": "test_output", "finished": false}}}',
          '{"serverContent": {"turnComplete": true}}',
      ]
  )
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  messages = session.receive()
  messages = await _async_iterator_to_list(messages)
  assert isinstance(messages[0], types.LiveServerMessage)
  assert messages[0].server_content.input_transcription.text == 'test_input'
  assert messages[0].server_content.input_transcription.finished == True

  assert isinstance(messages[1], types.LiveServerMessage)
  assert messages[1].server_content.output_transcription.text == 'test_output'
  assert messages[1].server_content.output_transcription.finished == False


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_go_away(
    mock_websocket, vertexai
):
  mock_websocket.recv = AsyncMock(
      side_effect=[
          '{"goAway": {"timeLeft": "10s"}}',
          '{"serverContent": {"turnComplete": true}}',
      ]
  )
  expected_result = types.LiveServerMessage(
      go_away=types.LiveServerGoAway(time_left='10s'),
  )
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  messages = session.receive()
  messages = await _async_iterator_to_list(messages)
  message = messages[0]

  assert isinstance(message, types.LiveServerMessage)
  assert message == expected_result


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_resumption_update(
    mock_websocket, vertexai
):
  mock_websocket.recv = AsyncMock(
      side_effect=[
          """{
                "sessionResumptionUpdate": {
                    "newHandle": "test_handle",
                    "resumable": "true",
                    "lastConsumedClientMessageIndex": "123456789"
                }
          }""",
          '{"serverContent": {"turnComplete": true}}',
      ]
  )

  expected_result = types.LiveServerMessage(
      session_resumption_update=types.LiveServerSessionResumptionUpdate(
          new_handle='test_handle',
          resumable=True,
          last_consumed_client_message_index=123456789
      ),
  )

  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  messages = session.receive()
  messages = await _async_iterator_to_list(messages)
  message = messages[0]

  assert isinstance(message, types.LiveServerMessage)
  assert message == expected_result


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_start_stream(
     mock_websocket, vertexai
):

  session = live.AsyncSession(
      mock_api_client(vertexai=vertexai), mock_websocket
  )

  async def mock_stream():
    yield b'data1'
    yield b'data2'

  async for message in session.start_stream(
      stream=mock_stream(), mime_type='audio/pcm'
  ):
    assert isinstance(message, types.LiveServerMessage)


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_receive_vad_signal(mock_websocket, vertexai):
  # Simulate the server sending a VAD signal message
  mock_websocket.recv = mock.AsyncMock(
      side_effect=[
          '{"voiceActivityDetectionSignal": {"vadSignalType": "VAD_SIGNAL_TYPE_SOS"}}',
          '{"serverContent": {"turnComplete": true}}',  # To close the receiver loop
      ]
  )

  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )

  messages = await _async_iterator_to_list(session.receive())

  # Check if the first message contains the VAD signal
  assert len(messages) > 0
  vad_message = messages[0]
  assert isinstance(vad_message, types.LiveServerMessage)
  assert vad_message.voice_activity_detection_signal is not None
  assert (
      vad_message.voice_activity_detection_signal.vad_signal_type
      == types.VadSignalType.VAD_SIGNAL_TYPE_SOS
  )

  # Check that the session can close cleanly
  assert messages[-1].server_content.turn_complete is True


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_close( mock_websocket, vertexai):
  session = live.AsyncSession(
      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_bidi_setup_to_api_no_config(vertexai):
  with warnings.catch_warnings():
    # Make sure there are no warnings cause by default values.
    warnings.simplefilter('error')
    result = await get_connect_message(
        mock_api_client(vertexai=vertexai),
        model='test_model'
    )
  expected_result = {'setup': {}}
  if vertexai:
    expected_result['setup']['model'] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
    expected_result['setup']['generationConfig'] = {}
    expected_result['setup']['generationConfig']['responseModalities'] = ["AUDIO"]
  else:
    expected_result['setup']['model'] = 'models/test_model'
  assert result == expected_result


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_speech_config(vertexai):

  expected_result = {
      'setup': {
          'model': 'models/test_model',
          'generationConfig': {
              'speechConfig': {
                  'voice_config': {
                      'prebuilt_voice_config': {'voice_name': 'en-default'}
                  },
                  'language_code': 'en-US',
              },
              'enableAffectiveDialog': True,
              'temperature': 0.7,
              'topP': 0.8,
              'topK': 9.0,
              'maxOutputTokens': 10,
              'mediaResolution': 'MEDIA_RESOLUTION_MEDIUM',
              'seed': 13,
          },
          'proactivity': {'proactive_audio': True},
          'systemInstruction': {
              'parts': [
                  {
                      'text': 'test instruction',
                  },
              ],
              'role': 'user',
          },
      }
  }
  if vertexai:
    expected_result['setup']['model'] = (
        'projects/test_project/locations/us-central1/'
        'publishers/google/models/test_model'
    )
    expected_result['setup']['generationConfig']['responseModalities'] = [
        'AUDIO'
    ]
    expected_result['setup']['generationConfig']['speechConfig'] = {
        'voiceConfig': {
            'prebuilt_voice_config': {'voice_name': 'en-default'}
        },
        'languageCode': 'en-US',
    }
  else:
    expected_result['setup']['model'] = 'models/test_model'

  # Config is a dict
  config_dict = {
      'speech_config': {
          'voice_config': {
              'prebuilt_voice_config': {'voice_name': 'en-default'}
          },
          'language_code': 'en-US',
      },
      'enable_affective_dialog': True,
      'proactivity': {'proactive_audio': True},
      'temperature': 0.7,
      'top_p': 0.8,
      'top_k': 9,
      'max_output_tokens': 10,
      'seed': 13,
      'system_instruction': 'test instruction',
      'media_resolution': 'MEDIA_RESOLUTION_MEDIUM',
  }
  result = await get_connect_message(
      mock_api_client(vertexai=vertexai), model='test_model', config=config_dict
  )
  assert types.LiveClientMessage._from_response(
      response=result, kwargs=None
  ) == types.LiveClientMessage._from_response(
      response=expected_result, kwargs=None
  )
  # Config is a LiveConnectConfig
  config = types.LiveConnectConfig(
      speech_config=types.SpeechConfig(
          voice_config=types.VoiceConfig(
              prebuilt_voice_config=types.PrebuiltVoiceConfig(
                  voice_name='en-default'
              )
          ),
          language_code='en-US',
      ),
      enable_affective_dialog=True,
      proactivity=types.ProactivityConfig(proactive_audio=True),
      temperature=0.7,
      top_p=0.8,
      top_k=9,
      max_output_tokens=10,
      media_resolution=types.MediaResolution.MEDIA_RESOLUTION_MEDIUM,
      seed=13,
      system_instruction='test instruction',
  )
  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config
  )
  assert types.LiveClientMessage._from_response(
      response=result, kwargs=None
  ) == types.LiveClientMessage._from_response(
      response=expected_result, kwargs=None
  )


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_error_if_multispeaker_voice_config(vertexai):

  # Config is a dict
  config_dict = {
      'speech_config': {
          'multi_speaker_voice_config': {
              'speaker_voice_configs': [
                  {
                      'speaker': 'Alice',
                      'voice_config': {
                          'prebuilt_voice_config': {'voice_name': 'leda'}
                      },
                  },
                  {
                      'speaker': 'Bob',
                      'voice_config': {
                          'prebuilt_voice_config': {'voice_name': 'kore'}
                      },
                  },
              ],
          },
      },
      'temperature': 0.7,
      'top_p': 0.8,
      'top_k': 9,
      'max_output_tokens': 10,
      'seed': 13,
      'system_instruction': 'test instruction',
      'media_resolution': 'MEDIA_RESOLUTION_MEDIUM',
  }
  with pytest.raises(ValueError, match='.*multi_speaker_voice_config.*'):
    result = await get_connect_message(
        mock_api_client(vertexai=vertexai),
        model='test_model',
        config=config_dict,
    )


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_explicit_vad(vertexai):
  # Config is a dict
  config_dict = {'explicit_vad_signal': True}
  with pytest_helper.exception_if_mldev(
      mock_api_client(vertexai=vertexai), ValueError
  ):
    result = await get_connect_message(
        mock_api_client(vertexai=vertexai), model='test_model', config=config_dict
    )
  if not vertexai:
    return
  assert result['setup']['explicitVadSignal'] == True


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_explicit_vad_config(vertexai):
  api_client = mock_api_client(vertexai=vertexai)

  # Config is a dict
  config_dict = {'explicit_vad_signal': True}
  with pytest_helper.exception_if_mldev(api_client, ValueError):
    result = await get_connect_message(
        mock_api_client(vertexai=vertexai),
        model='test_model',
        config=config_dict,
    )
  if not vertexai:
    return
  assert result['setup']['explicitVadSignal'] == True


@pytest.mark.asyncio
async def test_history_config():
  config_dict = {
      'history_config': {'initial_history_in_client_content': True},
  }
  result = await get_connect_message(
      mock_api_client(vertexai=False), model='test_model', config=config_dict
  )
  assert result['setup']['historyConfig'] == {
      'initial_history_in_client_content': True,
  }


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_system_instruction_as_content_type(
    vertexai,
):
  config_dict = {
      'system_instruction': {
          'parts': [{'text': 'test instruction'}],
          'role': 'user',
      },
  }
  config = types.LiveConnectConfig(**config_dict)
  expected_result = {
      'setup': {
          'model': 'test_model',
          'systemInstruction': {
              'parts': [{'text': 'test instruction'}],
              'role': 'user',
          },
      }
  }
  if vertexai:
    expected_result['setup'][
        'model'
    ] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
    expected_result['setup']['generationConfig'] = {}
    expected_result['setup']['generationConfig']['responseModalities'] = [
        'AUDIO'
    ]
  else:
    expected_result['setup']['model'] = 'models/test_model'

  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config
  )
  assert result == expected_result


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_config_tools_google_search(vertexai):
  config_dict = {
      'response_modalities': ['TEXT'],
      'system_instruction': 'test instruction',
      'generation_config': {'temperature': 0.7},
      'tools': [{'google_search': {}}],
  }

  config = types.LiveConnectConfig(**config_dict)
  expected_result = {
      'setup': {
          'generationConfig': {
              'temperature': 0.7,
              'responseModalities': ['TEXT'],
          },
          'systemInstruction': {
              'parts': [{'text': 'test instruction'}],
              'role': 'user',
          },
          'tools': [{'googleSearch': {}}],
      }
  }
  if vertexai:
    expected_result['setup']['model'] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
  else:
    expected_result['setup']['model'] = 'models/test_model'

  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config_dict
  )

  assert result == expected_result

  # Test config is a LiveConnectConfig
  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config
  )

  assert result == expected_result


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_config_tools_with_no_mcp(vertexai):
  config_dict = {
      'response_modalities': ['TEXT'],
      'system_instruction': 'test instruction',
      'generation_config': {'temperature': 0.7},
      'tools': [{'google_search': {}}],
  }

  config = types.LiveConnectConfig(**config_dict)
  expected_result = {
      'setup': {
          'generationConfig': {
              'temperature': 0.7,
              'responseModalities': ['TEXT'],
          },
          'systemInstruction': {
              'parts': [{'text': 'test instruction'}],
              'role': 'user',
          },
          'tools': [{'googleSearch': {}}],
      }
  }
  if vertexai:
    expected_result['setup']['model'] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
  else:
    expected_result['setup']['model'] = 'models/test_model'

  @patch.object(live, "McpClientSession", new=None)
  @patch.object(live, "McpTool", new=None)
  async def get_connect_message_no_mcp(config):
    return await get_connect_message(
        mock_api_client(vertexai=vertexai),
        model='test_model', config=config
    )

  result = await get_connect_message_no_mcp(config_dict)
  assert result == expected_result

  result = await get_connect_message_no_mcp(config_dict)
  assert result == expected_result


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_context_window_compression(
     vertexai
):
  config = types.LiveConnectConfig(
      generation_config=types.GenerationConfig(temperature=0.7),
      response_modalities=['TEXT'],
      system_instruction=types.Content(
          parts=[types.Part(text='test instruction')], role='user'
      ),
      context_window_compression=types.ContextWindowCompressionConfig(
          trigger_tokens=1000,
          sliding_window=types.SlidingWindow(target_tokens=10),
      ),
  )
  expected_result = {
      'setup': {
          'generationConfig': {
              'temperature': 0.7,
              'responseModalities': ['TEXT'],
          },
          'systemInstruction': {
              'parts': [{'text': 'test instruction'}],
              'role': 'user',
          },
           'contextWindowCompression': {
              'trigger_tokens': 1000,
              'sliding_window': {'target_tokens': 10},
          }
      }
  }
  if vertexai:
    expected_result['setup']['model'] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
  else:
    expected_result['setup']['model'] = 'models/test_model'

  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config
  )
  assert result == expected_result


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_config_tools_function_declaration(
     vertexai
):
  config_dict = {
      'generation_config': {'temperature': 0.7},
      'tools': [{'function_declarations': function_declarations}],
  }
  config = types.LiveConnectConfig(**config_dict)
  expected_result = {
      'setup': {
          'model': 'test_model',
          'tools': [{
              'functionDeclarations': [{
                  'parameters': {
                      'type': 'OBJECT',
                      'properties': {
                          'location': {
                              'type': 'STRING',
                              'description': (
                                  'The location to get the weather for'
                              ),
                          },
                          'unit': {'type': 'STRING', 'enum': ['C', 'F']},
                      },
                  },
                  'name': 'get_current_weather',
                  'description': 'Get the current weather in a city',
              }],
          }],
      }
  }
  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config
  )

  assert result['setup']['tools'][0]['functionDeclarations'][0][
      'description'
  ] == (
      expected_result['setup']['tools'][0]['functionDeclarations'][0][
          'description'
      ]
  )

  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config
  )

  assert result['setup']['tools'][0]['functionDeclarations'][0][
      'description'
  ] == (
      expected_result['setup']['tools'][0]['functionDeclarations'][0][
          'description'
      ]
  )


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_config_tools_function_directly(
     vertexai
):
  config_dict = {
      'generation_config': {'temperature': 0.7},
      'tools': [get_current_weather],
  }
  config = types.LiveConnectConfig(**config_dict)
  expected_result = {
      'setup': {
          'model': 'test_model',
          'tools': [{
              'functionDeclarations': [{
                  'parameters': {
                      'type': 'OBJECT',
                      'properties': {
                          'location': {
                              'type': 'STRING',
                              'description': (
                                  'The location to get the weather for'
                              ),
                          },
                          'unit': {'type': 'STRING', 'enum': ['C', 'F']},
                      },
                  },
                  'name': 'get_current_weather',
                  'description': 'Get the current weather in a city.',
              }],
          }],
      }
  }
  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config
  )

  assert result['setup']['tools'][0]['functionDeclarations'][0][
      'description'
  ] == (
      expected_result['setup']['tools'][0]['functionDeclarations'][0][
          'description'
      ]
  )

  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config
  )

  assert result['setup']['tools'][0]['functionDeclarations'][0][
      'description'
  ] == (
      expected_result['setup']['tools'][0]['functionDeclarations'][0][
          'description'
      ]
  )


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_tools_function_behavior(vertexai):
  api_client = mock_api_client(vertexai=vertexai)

  declaration = types.FunctionDeclaration.from_callable(
      client=api_client, callable=get_current_weather
  )
  declaration.behavior = types.Behavior.NON_BLOCKING
  config_dict = {
      'generation_config': {'temperature': 0.7},
      'tools': [{'function_declarations': [declaration]}],
  }
  config = types.LiveConnectConfig(**config_dict)

  with pytest_helper.exception_if_vertex(api_client, ValueError):
    result = await get_connect_message(
        mock_api_client(vertexai=vertexai), model='test_model', config=config
    )
  if vertexai:
    return

  assert (
      result['setup']['tools'][0]['functionDeclarations'][0]['behavior']
      == 'NON_BLOCKING'
  )


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_config_mcp_tools(
    vertexai,
):
  if mcp_types is None:
    return

  expected_result_googleai = {
      'setup': {
          'model': 'models/test_model',
          'tools': [{
              'functionDeclarations': [{
                  'parameters': {
                      'type': 'OBJECT',
                      'properties': {
                          'location': {
                              'type': 'STRING',
                          },
                      },
                  },
                  'name': 'get_weather',
                  'description': 'Get the weather in a city.',
              }],
          }],
      }
  }
  expected_result_vertexai = {
      'setup': {
          'generationConfig': {
              'responseModalities': [
                  'AUDIO',
              ],
          },
          'model': (
              'projects/test_project/locations/us-central1/publishers/google/models/test_model'
          ),
          'tools': [{
              'functionDeclarations': [{
                  'parameters': {
                      'type': 'OBJECT',
                      'properties': {
                          'location': {
                              'type': 'STRING',
                          },
                      },
                  },
                  'name': 'get_weather',
                  'description': 'Get the weather in a city.',
              }],
          }],
      }
  }
  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model',
      config={
          'tools': [
              mcp_types.Tool(
                  name='get_weather',
                  description='Get the weather in a city.',
                  inputSchema={
                      'type': 'object',
                      'properties': {'location': {'type': 'string'}},
                  },
              )
          ],
      },
  )

  assert (
      result == expected_result_vertexai
      if vertexai
      else expected_result_googleai
  )


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_config_mcp_session(
    vertexai,
):
  if mcp_types is None:
    return

  class MockMcpClientSession(McpClientSession):

    def __init__(self):
      self._read_stream = None
      self._write_stream = None

    async def list_tools(self):
      return mcp_types.ListToolsResult(
          tools=[
              mcp_types.Tool(
                  name='get_weather',
                  description='Get the weather in a city.',
                  inputSchema={
                      'type': 'object',
                      'properties': {'location': {'type': 'string'}},
                  },
              ),
          ]
      )

  expected_result_googleai = {
      'setup': {
          'model': 'models/test_model',
          'tools': [{
              'functionDeclarations': [{
                  'parameters': {
                      'type': 'OBJECT',
                      'properties': {
                          'location': {
                              'type': 'STRING',
                          },
                      },
                  },
                  'name': 'get_weather',
                  'description': 'Get the weather in a city.',
              }],
          }],
      }
  }
  expected_result_vertexai = {
      'setup': {
          'generationConfig': {
              'responseModalities': [
                  'AUDIO',
              ],
          },
          'model': (
              'projects/test_project/locations/us-central1/publishers/google/models/test_model'
          ),
          'tools': [{
              'functionDeclarations': [{
                  'parameters': {
                      'type': 'OBJECT',
                      'properties': {
                          'location': {
                              'type': 'STRING',
                          },
                      },
                  },
                  'name': 'get_weather',
                  'description': 'Get the weather in a city.',
              }],
          }],
      }
  }
  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model',
      config={
          'tools': [MockMcpClientSession()],
      },
  )

  assert (
      result == expected_result_vertexai
      if vertexai
      else expected_result_googleai
  )


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_config_tools_code_execution(
     vertexai
):
  config_dict = {
      'tools': [{'code_execution': {}}],
  }
  config = types.LiveConnectConfig(**config_dict)
  expected_result = {
      'setup': {
          'model': 'test_model',
          'tools': [{
              'codeExecution': {},
          }],
      }
  }
  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config
  )

  assert result['setup']['tools'][0] == expected_result['setup']['tools'][0]


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_realtime_input_config(vertexai):
  config_dict = {
      'realtime_input_config': {
          'automatic_activity_detection': {
              'disabled': True,
              'start_of_speech_sensitivity': 'START_SENSITIVITY_HIGH',
              'end_of_speech_sensitivity': 'END_SENSITIVITY_HIGH',
              'prefix_padding_ms': 20,
              'silence_duration_ms': 100,
          },
          'activity_handling': 'NO_INTERRUPTION',
          'turn_coverage': 'TURN_INCLUDES_ALL_INPUT',
      }
  }

  config = types.LiveConnectConfig(**config_dict)
  expected_result = {
      'setup': {
          'model': 'test_model',
          'realtimeInputConfig': config_dict['realtime_input_config'],
      }
  }

  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config
  )

  assert (
      result['setup']['realtimeInputConfig']
      == expected_result['setup']['realtimeInputConfig']
  )


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_input_transcription(vertexai):
  config_dict = {
      'input_audio_transcription': {},
  }
  config = types.LiveConnectConfig(**config_dict)
  expected_result = {
      'setup': {
          'model': 'test_model',
          'inputAudioTranscription': {},
      }
  }

  result = await get_connect_message(
      mock_api_client(vertexai=vertexai), model='test_model', config=config
  )

  assert (
      result['setup']['inputAudioTranscription']
      == expected_result['setup']['inputAudioTranscription']
  )


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_output_transcription(vertexai):
  config_dict = {
      'output_audio_transcription': {},
  }
  config = types.LiveConnectConfig(**config_dict)
  expected_result = {
      'setup': {
          'model': 'test_model',
          'outputAudioTranscription': {},
      }
  }

  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config
  )

  assert (
      result['setup']['outputAudioTranscription']
      == expected_result['setup']['outputAudioTranscription']
  )

@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_media_resolution(vertexai):
  config_dict = {
      'media_resolution': 'MEDIA_RESOLUTION_LOW',
  }
  config = types.LiveConnectConfig(**config_dict)
  expected_result = {
      'setup': {
          'model': 'test_model',
          'generationConfig': {'mediaResolution':'MEDIA_RESOLUTION_LOW'},
      }
  }

  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model', config=config
  )

  assert (
      result['setup']['generationConfig']['mediaResolution']
      == expected_result['setup']['generationConfig']['mediaResolution']
  )


@pytest.mark.parametrize('vertexai', [True])
@pytest.mark.asyncio
async def test_bidi_setup_publishers(
     vertexai
):
  expected_result = {
      'setup': {
         'generationConfig': {
             'responseModalities': [
                 'AUDIO',
             ],
         },
         'model': 'projects/test_project/locations/us-central1/publishers/google/models/test_model',
      }
  }
  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='publishers/google/models/test_model')

  assert result == expected_result


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_generation_config_warning(
     vertexai
):
  with pytest.warns(
      DeprecationWarning,
      match='Setting `LiveConnectConfig.generation_config` is deprecated'
  ):
    result = await get_connect_message(
        mock_api_client(vertexai=vertexai),
        model='models/test_model',
        config={'generation_config': {'temperature': 0.7}})

  assert result['setup']['generationConfig']['temperature'] == 0.7


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_session_resumption(vertexai):
  config_dict = {
      'session_resumption': {'handle': 'test_handle'},
  }
  config = types.LiveConnectConfig(**config_dict)

  result = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model',
      config=config
  )
  expected_result = {
      'setup': {
          'sessionResumption': {
              'handle': 'test_handle',
          },
      }
  }
  if vertexai:
    expected_result['setup']['generationConfig'] = {
        'responseModalities': [
            'AUDIO',
        ],
    }
    expected_result['setup']['model'] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
  else:
    expected_result['setup']['model'] = 'models/test_model'
  assert result == expected_result


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_transparent_session_resumption(vertexai):
  api_client = mock_api_client(vertexai=vertexai)
  config_dict = {
      'session_resumption': {'handle': 'test_handle', 'transparent': True},
  }
  config = types.LiveConnectConfig(**config_dict)

  with pytest_helper.exception_if_mldev(api_client, ValueError):
    result = await get_connect_message(
        mock_api_client(vertexai=vertexai),
        model='test_model',
        config=config
    )

  expected_result = {
      'setup': {
          'sessionResumption': {
              'handle': 'test_handle',
              'transparent': True,
          },
      }
  }
  if vertexai:
    expected_result['setup']['generationConfig'] = {
        'responseModalities': [
            'AUDIO',
        ],
    }
    expected_result['setup']['model'] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
  else:
    return

  assert result == expected_result


@pytest.mark.parametrize('vertexai', [True, False])
def test_parse_client_message_str( mock_websocket, vertexai):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  result = session._parse_client_message('test')
  assert 'client_content' in result
  assert result == {
      'client_content': {
          'turn_complete': False,
          'turns': [{'role': 'user', 'parts': [{'text': 'test'}]}],
      }
  }
  # _parse_client_message returns a TypedDict, so we should be able to
  # construct a LiveClientMessage from it
  assert types.LiveClientMessage(**result)


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_thinking_config(vertexai):
  config_dict = {
      'thinking_config': {
          'include_thoughts': True,
          'thinking_budget': 1024,
      }
  }

  expected_gen_config = {
      'thinkingConfig': {
          'include_thoughts': True,
          'thinking_budget': 1024,
      }
  }

  if vertexai:
    expected_gen_config['responseModalities'] = ['AUDIO']

  expected_result = {
      'setup': {
          'generationConfig': expected_gen_config,
      }
  }

  if vertexai:
    expected_result['setup'][
        'model'
    ] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
  else:
    expected_result['setup']['model'] = 'models/test_model'

  result = await get_connect_message(
      mock_api_client(vertexai=vertexai), model='test_model', config=config_dict
  )
  result = pytest_helper.camel_to_snake_all_keys(result)
  expected_result = pytest_helper.camel_to_snake_all_keys(expected_result)
  assert result == expected_result


@pytest.mark.parametrize('vertexai', [True, False])
def test_parse_client_message_blob( mock_websocket, vertexai):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  result = session._parse_client_message(
      types.Blob(data=bytes([0, 0, 0]), mime_type='text/plain')
  )
  assert 'realtime_input' in result
  assert result == {
      'realtime_input': {
          'media_chunks': [{'mime_type': 'text/plain', 'data': 'AAAA'}],
      }
  }


@pytest.mark.parametrize('vertexai', [True, False])
def test_parse_client_message_blob_dict(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )

  blob = types.Blob(data=bytes([0, 0, 0]), mime_type='text/plain')
  blob_dict = blob.model_dump()
  result = session._parse_client_message(blob_dict)
  assert 'realtime_input' in result
  assert result == {
      'realtime_input': {
          'media_chunks': [{'mime_type': 'text/plain', 'data': 'AAAA'}],
      }
  }


@pytest.mark.parametrize('vertexai', [True, False])
def test_parse_client_message_client_content(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  result = session._parse_client_message(
      types.LiveClientContent(
          turn_complete=False,
          turns=[types.Content(parts=[types.Part(text='test')], role='user')],
      )
  )
  assert 'client_content' in result
  assert result == {
      'client_content': {
          'turn_complete': False,
          'turns': [{'role': 'user', 'parts': [{'text': 'test'}]}],
      }
  }


@pytest.mark.parametrize('vertexai', [True, False])
def test_parse_client_message_client_content_blob(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  client_content = types.LiveClientContent(
      turn_complete=False,
      turns=[
          types.Content(
              parts=[
                  types.Part(
                      inline_data=types.Blob(
                          data=bytes([0, 0, 0]), mime_type='text/plain'
                      )
                  )
              ],
              role='user',
          )
      ],
  )
  result = session._parse_client_message(client_content)
  assert 'client_content' in result
  assert (
      type(
          result['client_content']['turns'][0]['parts'][0]['inline_data'][
              'data'
          ]
      )
      == str
  )
  assert result == {
      'client_content': {
          'turn_complete': False,
          'turns': [{
              'role': 'user',
              'parts': [
                  {'inline_data': {'mime_type': 'text/plain', 'data': 'AAAA'}}
              ],
          }],
      }
  }


@pytest.mark.parametrize('vertexai', [True, False])
def test_parse_client_message_client_content_dict(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  client_content = types.LiveClientContent(
      turn_complete=False,
      turns=[
          types.Content(
              parts=[
                  types.Part(
                      inline_data=types.Blob(
                          data=bytes([0, 0, 0]), mime_type='text/plain'
                      )
                  )
              ],
              role='user',
          )
      ],
  )
  result = session._parse_client_message(
      client_content.model_dump(mode='json', exclude_none=True)
  )
  assert 'client_content' in result
  assert (
      type(
          result['client_content']['turns'][0]['parts'][0]['inline_data'][
              'data'
          ]
      )
      == str
  )
  assert result == {
      'client_content': {
          'turn_complete': False,
          'turns': [{
              'role': 'user',
              'parts': [
                  {'inline_data': {'mime_type': 'text/plain', 'data': 'AAAA'}}
              ],
          }],
      }
  }


@pytest.mark.parametrize('vertexai', [True, False])
def test_parse_client_message_realtime_input(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  input = types.LiveClientRealtimeInput(
      media_chunks=[types.Blob(data=bytes([0, 0, 0]), mime_type='text/plain')]
  )
  result = session._parse_client_message(input)
  assert 'realtime_input' in result
  assert result == {
      'realtime_input': {
          'media_chunks': [{'mime_type': 'text/plain', 'data': 'AAAA'}],
      }
  }


@pytest.mark.parametrize('vertexai', [True, False])
def test_parse_client_message_realtime_input_dict(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  input = types.LiveClientRealtimeInput(
      media_chunks=[types.Blob(data=bytes([0, 0, 0]), mime_type='text/plain')]
  )
  result = session._parse_client_message(
      input.model_dump(mode='json', exclude_none=True)
  )
  assert 'realtime_input' in result
  assert result == {
      'realtime_input': {
          'media_chunks': [{'mime_type': 'text/plain', 'data': 'AAAA'}],
      }
  }


@pytest.mark.parametrize('vertexai', [True, False])
def test_parse_client_message_tool_response(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  input = types.LiveClientToolResponse(
      function_responses=[
          types.FunctionResponse(
              id='test_id',
              name='test_name',
              response={'result': 'test_response'},
          )
      ]
  )
  result = session._parse_client_message(input)
  assert 'tool_response' in result
  assert result == {
      'tool_response': {
          'function_responses': [
              {
                  'id': 'test_id',
                  'name': 'test_name',
                  'response': {
                      'result': 'test_response',
                  },
              },
          ],
      }
  }


@pytest.mark.parametrize('vertexai', [True, False])
def test_parse_client_message_function_response(
    mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  input = types.FunctionResponse(
    id='test_id',
    name='test_name',
    response={
        'result': 'test_response',
        'user_name': 'test_user_name',
        'userEmail': 'test_user_email',
    },
  )
  result = session._parse_client_message(input)
  assert 'tool_response' in result
  assert result == {
      'tool_response': {
          'function_responses': [
              {
                  'id': 'test_id',
                  'name': 'test_name',
                  'response': {
                      'result': 'test_response',
                      'user_name': 'test_user_name',
                      'userEmail': 'test_user_email',
                  },
              },
          ],
      }
  }


@pytest.mark.parametrize('vertexai', [True, False])
def test_parse_client_message_tool_response_dict_with_only_response(
     mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  input = {
    'id': 'test_id',
    'name': 'test_name',
    'response': {
        'result': 'test_response',
    }
  }
  result = session._parse_client_message(input)
  assert 'tool_response' in result
  assert result == {
      'tool_response': {
          'function_responses': [
              {
                  'id': 'test_id',
                  'name': 'test_name',
                  'response': {
                      'result': 'test_response',
                  },
              },
          ],
      }
  }


@pytest.mark.parametrize('vertexai', [True, False])
def test_parse_client_message_realtime_tool_response(
    mock_websocket, vertexai
):
  session = live.AsyncSession(
      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
  )
  input = types.LiveClientToolResponse(
      function_responses=[
          types.FunctionResponse(
              id='test_id',
              name='test_name',
              response={'result': 'test_response'},
          )
      ]
  )

  result = session._parse_client_message(
      input.model_dump(mode='json', exclude_none=True)
  )
  assert 'tool_response' in result
  assert result == {
      'tool_response': {
          'function_responses': [
              {
                  'id': 'test_id',
                  'name': 'test_name',
                  'response': {
                      'result': 'test_response',
                  },
              },
          ],
      }
  }


@pytest.mark.asyncio
async def test_connect_with_provided_credentials(mock_websocket):
  # custom oauth2 credentials
  credentials = Credentials(token='provided_fake_token')
  # mock api client
  client = mock_api_client(vertexai=True, credentials=credentials)
  capture = {}

  @contextlib.asynccontextmanager
  async def mock_connect(uri, additional_headers=None, **kwargs):
    capture['headers'] = additional_headers
    yield mock_websocket

  @patch.object(live, 'ws_connect', new=mock_connect)
  async def _test_connect():
    live_module = live.AsyncLive(client)
    async with live_module.connect(model='test-model'):
      pass

    assert 'Authorization' in capture['headers']
    assert capture['headers']['Authorization'] == 'Bearer provided_fake_token'

  await _test_connect()


@pytest.mark.asyncio
async def test_connect_with_default_credentials(mock_websocket):
  # mock api client
  client = mock_api_client(vertexai=True, credentials=None)
  # mock google auth cred
  mock_google_auth_default = Mock(return_value=(None, None))
  mock_creds = Mock(token='default_test_token')
  mock_google_auth_default.return_value = (mock_creds, None)
  capture = {}

  @contextlib.asynccontextmanager
  async def mock_connect(uri, additional_headers=None, **kwargs):
    capture['headers'] = additional_headers
    yield mock_websocket

  @patch('google.auth.default', new=mock_google_auth_default)
  @patch.object(live, 'ws_connect', new=mock_connect)
  async def _test_connect():
    live_module = live.AsyncLive(client)
    async with live_module.connect(model='test-model'):
      pass

    assert 'Authorization' in capture['headers']
    assert capture['headers']['Authorization'] == 'Bearer default_test_token'

  await _test_connect()


@pytest.mark.asyncio
async def test_connect_with_custom_base_url(mock_websocket):
  # mock api client
  client = gl_client.BaseApiClient(
      vertexai=True,
      http_options={
          'base_url': 'https://custom-base-url.com',
          'headers': {'Authorization': 'Bearer custom_test_token'},
      }
  )
  # No ADC credentials.
  capture = {}

  @contextlib.asynccontextmanager
  async def mock_connect(uri, additional_headers=None, **kwargs):
    capture['uri'] = uri
    capture['headers'] = additional_headers
    yield mock_websocket

  @patch.object(live, 'ws_connect', new=mock_connect)
  async def _test_connect():
    live_module = live.AsyncLive(client)
    async with live_module.connect(model='test-model'):
      pass

    assert 'Authorization' in capture['headers']
    assert capture['headers']['Authorization'] == 'Bearer custom_test_token'
    assert capture['uri'] == 'https://custom-base-url.com'

  await _test_connect()


@pytest.mark.parametrize('vertexai', [False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_auth_tokens(mock_websocket, vertexai):
  api_client_mock = mock_api_client(vertexai=vertexai)
  api_client_mock.api_key = 'auth_tokens/TEST_AUTH_TOKEN'
  result = await get_connect_message(api_client_mock, model='test_model')

  mock_ws = AsyncMock()
  mock_ws.send = AsyncMock()
  mock_ws.recv = AsyncMock(
    return_value=(
        b'{\n  "setupComplete": {"sessionId": "test_session_id"}\n}\n'
    )
  )
  capture = {}

  @contextlib.asynccontextmanager
  async def mock_connect(uri, additional_headers=None, **kwargs):
    capture['uri'] = uri
    capture['headers'] = additional_headers
    yield mock_ws

  with patch.object(live, 'ws_connect', new=mock_connect):
    live_module = live.AsyncLive(api_client_mock)
    async with live_module.connect(
        model='test_model',
    ):
      pass

  assert (
      'Authorization' in capture['headers']
  ), 'Authorization key is missing from headers'
  assert (
      capture['headers']['Authorization'] == 'Token auth_tokens/TEST_AUTH_TOKEN'
  )
  assert 'BidiGenerateContentConstrained' in capture['uri']


@pytest.mark.parametrize('vertexai', [False])
@pytest.mark.asyncio
async def test_bidi_setup_to_api_with_api_key(mock_websocket, vertexai):
  api_client_mock = mock_api_client(vertexai=vertexai)
  api_client_mock._http_options = types.HttpOptions.model_validate(
      {'headers': {'x-goog-api-key': 'TEST_API_KEY'}}
  )
  result = await get_connect_message(api_client_mock, model='test_model')

  mock_ws = AsyncMock()
  mock_ws.send = AsyncMock()
  mock_ws.recv = AsyncMock(
    return_value=(
        b'{\n  "setupComplete": {"sessionId": "test_session_id"}\n}\n'
    )
  )
  capture = {}

  @contextlib.asynccontextmanager
  async def mock_connect(uri, additional_headers=None, **kwargs):
    capture['uri'] = uri
    capture['headers'] = additional_headers
    yield mock_ws

  with patch.object(live, 'ws_connect', new=mock_connect):
    live_module = live.AsyncLive(api_client_mock)
    async with live_module.connect(
        model='test_model',
    ):
      pass

  assert 'x-goog-api-key' in capture['headers'], "x-goog-api-key is missing from headers"
  assert capture['headers']['x-goog-api-key'] == 'TEST_API_KEY'
  assert 'BidiGenerateContent' in capture['uri']


@pytest.mark.parametrize('vertexai', [True, False])
@pytest.mark.asyncio
async def test_async_session_setup_complete_with_voice_consent_signature(
    vertexai,
):
  mock_ws = AsyncMock()
  mock_ws.send = AsyncMock()
  mock_ws.recv = AsyncMock(
      return_value=(
          b'{"setupComplete": {"sessionId": "test_session_id",'
          b' "voiceConsentSignature": {"signature": "test_sig_abc123"}}}'
      )
  )
  mock_ws.close = AsyncMock()

  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, **kwargs):
    yield mock_ws

  @patch('google.auth.default', new=mock_google_auth_default)
  @patch.object(live, 'ws_connect', new=mock_connect)
  async def _test_connect():
    live_module = live.AsyncLive(mock_api_client(vertexai=vertexai))
    async with live_module.connect(model='test_model') as session:
      assert session.setup_complete is not None
      assert session.setup_complete.session_id == 'test_session_id'
      assert session.setup_complete.voice_consent_signature is not None
      assert (
          session.setup_complete.voice_consent_signature.signature
          == 'test_sig_abc123'
      )

  await _test_connect()


@pytest.mark.parametrize('vertexai', [False])
@pytest.mark.asyncio
async def test_bidi_setup_replicated_voice_config_with_consent(vertexai):
  config = types.LiveConnectConfig(
      response_modalities=['AUDIO'],
      speech_config=types.SpeechConfig(
          voice_config=types.VoiceConfig(
              replicated_voice_config=types.ReplicatedVoiceConfig(
                  mime_type='audio/wav',
                  voice_sample_audio=b'fake_audio_data',
                  consent_audio=b'fake_consent_data',
              )
          )
      ),
  )
  result = await get_connect_message(
      mock_api_client(vertexai=vertexai), model='test_model', config=config
  )

  setup = result.get('setup', {})
  gen_config = setup.get('generationConfig', {})
  speech_config = gen_config.get('speechConfig', {})
  voice_config = speech_config.get('voice_config', {})
  replicated = voice_config.get('replicated_voice_config', {})

  assert replicated.get('mime_type') == 'audio/wav'
  assert replicated.get('voice_sample_audio') is not None
  assert replicated.get('consent_audio') is not None

  config_with_sig = types.LiveConnectConfig(
      response_modalities=['AUDIO'],
      speech_config=types.SpeechConfig(
          voice_config=types.VoiceConfig(
              replicated_voice_config=types.ReplicatedVoiceConfig(
                  mime_type='audio/wav',
                  voice_sample_audio=b'fake_audio_data',
                  voice_consent_signature=types.VoiceConsentSignature(
                      signature='test_sig_abc123'
                  ),
              )
          )
      ),
  )
  result_with_sig = await get_connect_message(
      mock_api_client(vertexai=vertexai),
      model='test_model',
      config=config_with_sig,
  )

  setup_sig = result_with_sig.get('setup', {})
  gen_config_sig = setup_sig.get('generationConfig', {})
  speech_config_sig = gen_config_sig.get('speechConfig', {})
  voice_config_sig = speech_config_sig.get('voice_config', {})
  replicated_sig = voice_config_sig.get('replicated_voice_config', {})

  assert replicated_sig.get('mime_type') == 'audio/wav'
  assert replicated_sig.get('voice_sample_audio') is not None
  assert replicated_sig.get('voice_consent_signature') is not None
  assert (
      replicated_sig['voice_consent_signature'].get('signature')
      == 'test_sig_abc123'
  )

