"""Tests for Google provider functionality (both chat and ASR).""" import io import json from unittest.mock import MagicMock, mock_open, patch import pytest from aisuite.providers.google_provider import GoogleProvider from aisuite.provider import ASRError from aisuite.framework.message import ( TranscriptionResult, TranscriptionOptions, StreamingTranscriptionChunk, ) from vertexai.generative_models import Content, Part @pytest.fixture(autouse=True) def set_api_key_env_var(monkeypatch): """Fixture to set environment variables for tests.""" monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "path-to-service-account-json") monkeypatch.setenv("GOOGLE_PROJECT_ID", "vertex-project-id") monkeypatch.setenv("GOOGLE_REGION", "us-central1") @pytest.fixture def mock_google_speech_response(): """Create a mock Google Speech-to-Text API response.""" mock_response = MagicMock() mock_result = MagicMock() mock_alternative = MagicMock() mock_alternative.transcript = "Hello, this is a test transcription." mock_alternative.confidence = 0.95 # Mock words mock_word = MagicMock() mock_word.word = "Hello" mock_word.start_time.total_seconds.return_value = 0.0 mock_word.end_time.total_seconds.return_value = 0.5 mock_word.confidence = 0.98 mock_alternative.words = [mock_word] mock_result.alternatives = [mock_alternative] mock_response.results = [mock_result] return mock_response def test_missing_env_vars(): """Test that an error is raised if required environment variables are missing.""" with patch.dict("os.environ", {}, clear=True): with pytest.raises(EnvironmentError) as exc_info: GoogleProvider() assert "Missing one or more required Google environment variables" in str( exc_info.value ) def test_vertex_interface(): """High-level test that the interface is initialized and chat completions are requested successfully.""" # Test case 1: Regular text response def test_text_response(): user_greeting = "Hello!" message_history = [{"role": "user", "content": user_greeting}] selected_model = "our-favorite-model" response_text_content = "mocked-text-response-from-model" interface = GoogleProvider() mock_response = MagicMock() mock_response.candidates = [MagicMock()] mock_response.candidates[0].content.parts = [MagicMock()] mock_response.candidates[0].content.parts[0].text = response_text_content # Ensure function_call attribute doesn't exist del mock_response.candidates[0].content.parts[0].function_call with patch( "aisuite.providers.google_provider.GenerativeModel" ) as mock_generative_model: mock_model = MagicMock() mock_generative_model.return_value = mock_model mock_chat = MagicMock() mock_model.start_chat.return_value = mock_chat mock_chat.send_message.return_value = mock_response response = interface.chat_completions_create( messages=message_history, model=selected_model, temperature=0.7, ) # Assert the response is in the correct format assert response.choices[0].message.content == response_text_content assert response.choices[0].finish_reason == "stop" # Test case 2: Function call response def test_function_call(): user_greeting = "What's the weather?" message_history = [{"role": "user", "content": user_greeting}] selected_model = "our-favorite-model" interface = GoogleProvider() mock_response = MagicMock() mock_response.candidates = [MagicMock()] mock_response.candidates[0].content.parts = [MagicMock()] # Mock the function call response function_call_mock = MagicMock() function_call_mock.name = "get_weather" function_call_mock.args = {"location": "San Francisco"} mock_response.candidates[0].content.parts[0].function_call = function_call_mock mock_response.candidates[0].content.parts[0].text = None with patch( "aisuite.providers.google_provider.GenerativeModel" ) as mock_generative_model: mock_model = MagicMock() mock_generative_model.return_value = mock_model mock_chat = MagicMock() mock_model.start_chat.return_value = mock_chat mock_chat.send_message.return_value = mock_response response = interface.chat_completions_create( messages=message_history, model=selected_model, temperature=0.7, ) # Assert the response contains the function call assert response.choices[0].message.content is None assert response.choices[0].message.tool_calls[0].type == "function" assert ( response.choices[0].message.tool_calls[0].function.name == "get_weather" ) assert json.loads( response.choices[0].message.tool_calls[0].function.arguments ) == {"location": "San Francisco"} assert response.choices[0].finish_reason == "tool_calls" # Run both test cases test_text_response() test_function_call() def test_convert_openai_to_vertex_ai(): """Test the message conversion from OpenAI format to Vertex AI format.""" interface = GoogleProvider() message = {"role": "user", "content": "Hello!"} # Use the transformer to convert the message result = interface.transformer.convert_request([message]) # Verify the conversion result assert len(result) == 1 assert isinstance(result[0], Content) assert result[0].role == "user" assert len(result[0].parts) == 1 assert isinstance(result[0].parts[0], Part) assert result[0].parts[0].text == "Hello!" def test_role_conversions(): """Test that different message roles are converted correctly.""" interface = GoogleProvider() messages = [ {"role": "system", "content": "System message"}, {"role": "user", "content": "User message"}, {"role": "assistant", "content": "Assistant message"}, ] result = interface.transformer.convert_request(messages) # System and user messages should both be converted to "user" role in Vertex AI assert len(result) == 3 assert result[0].role == "user" # system converted to user assert result[0].parts[0].text == "System message" assert result[1].role == "user" assert result[1].parts[0].text == "User message" assert result[2].role == "model" # assistant converted to model assert result[2].parts[0].text == "Assistant message" class TestGoogleProvider: """Test suite for Google provider functionality.""" def test_provider_initialization(self): """Test that Google provider initializes correctly.""" provider = GoogleProvider() assert provider is not None assert hasattr(provider, "audio") assert hasattr(provider.audio, "transcriptions") class TestGoogleASR: """Test suite for Google ASR functionality.""" def test_audio_transcriptions_create_success(self, mock_google_speech_response): """Test successful audio transcription.""" google_provider = GoogleProvider() mock_client = MagicMock() mock_client.recognize.return_value = mock_google_speech_response google_provider._speech_client = mock_client with patch("builtins.open", mock_open(read_data=b"fake audio data")): result = google_provider.audio.transcriptions.create( model="google:latest_long", file="test_audio.wav" ) assert isinstance(result, TranscriptionResult) assert result.text == "Hello, this is a test transcription." assert result.confidence == 0.95 assert result.task == "transcribe" def test_audio_transcriptions_create_with_file_object( self, mock_google_speech_response ): """Test audio transcription with file-like object.""" google_provider = GoogleProvider() mock_client = MagicMock() mock_client.recognize.return_value = mock_google_speech_response google_provider._speech_client = mock_client audio_data = io.BytesIO(b"fake audio data") result = google_provider.audio.transcriptions.create( model="google:latest_long", file=audio_data ) assert isinstance(result, TranscriptionResult) assert result.text == "Hello, this is a test transcription." def test_audio_transcriptions_create_with_options( self, mock_google_speech_response ): """Test audio transcription with TranscriptionOptions.""" google_provider = GoogleProvider() mock_client = MagicMock() mock_client.recognize.return_value = mock_google_speech_response google_provider._speech_client = mock_client options = TranscriptionOptions( language="en", enable_automatic_punctuation=True, include_word_timestamps=True, ) with patch("builtins.open", mock_open(read_data=b"fake audio data")): result = google_provider.audio.transcriptions.create( model="google:latest_long", file="test_audio.wav", options=options ) mock_client.recognize.assert_called_once() assert isinstance(result, TranscriptionResult) assert result.text == "Hello, this is a test transcription." def test_audio_transcriptions_create_error_handling(self): """Test handling of Google Speech API errors.""" google_provider = GoogleProvider() mock_client = MagicMock() mock_client.recognize.side_effect = Exception("API Error") google_provider._speech_client = mock_client with patch("builtins.open", mock_open(read_data=b"fake audio data")): with pytest.raises( ASRError, match="Google Speech-to-Text error: API Error" ): google_provider.audio.transcriptions.create( model="google:latest_long", file="test_audio.wav" ) @pytest.mark.asyncio async def test_audio_transcriptions_create_stream_output( self, mock_google_speech_response ): """Test streaming audio transcription with fixed config parameter.""" google_provider = GoogleProvider() mock_client = MagicMock() # Mock streaming response mock_streaming_response = MagicMock() mock_streaming_result = MagicMock() mock_alternative = MagicMock() mock_alternative.transcript = "Hello streaming" mock_streaming_result.alternatives = [mock_alternative] mock_streaming_result.is_final = True mock_streaming_response.results = [mock_streaming_result] mock_client.streaming_recognize.return_value = [mock_streaming_response] google_provider._speech_client = mock_client with patch("builtins.open", mock_open(read_data=b"fake audio data")): result = google_provider.audio.transcriptions.create_stream_output( model="google:latest_long", file="test_audio.wav" ) assert hasattr(result, "__aiter__") # Test that streaming_recognize is called with both config and requests chunks = [] async for chunk in result: chunks.append(chunk) break # Just test one chunk # Verify the method was called with correct parameters mock_client.streaming_recognize.assert_called_once() call_args = mock_client.streaming_recognize.call_args assert "config" in call_args.kwargs assert "requests" in call_args.kwargs def test_parse_google_response_complete(self, mock_google_speech_response): """Test parsing complete Google response.""" google_provider = GoogleProvider() result = google_provider.audio.transcriptions._parse_google_response( mock_google_speech_response ) assert result.text == "Hello, this is a test transcription." assert result.confidence == 0.95 assert result.task == "transcribe" assert len(result.words) == 1 word = result.words[0] assert word.word == "Hello" assert word.start == 0.0 assert word.end == 0.5 assert word.confidence == 0.98 assert len(result.alternatives) == 1 def test_parse_google_response_empty_results(self): """Test parsing response with empty results.""" google_provider = GoogleProvider() mock_response = MagicMock() mock_response.results = [] result = google_provider.audio.transcriptions._parse_google_response( mock_response ) assert result.text == "" assert result.language is None