# Copyright 2026 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. from __future__ import annotations import os from unittest import mock from unittest.mock import AsyncMock from google.adk.models.apigee_llm import ApigeeLlm from google.adk.models.apigee_llm import CompletionsHTTPClient from google.adk.models.llm_request import LlmRequest from google.genai import types from google.genai.types import Content from google.genai.types import Part import pytest BASE_MODEL_ID = 'gemini-2.5-flash' APIGEE_GEMINI_MODEL_ID = 'apigee/gemini/v1/' + BASE_MODEL_ID APIGEE_VERTEX_MODEL_ID = 'apigee/vertex_ai/v1beta/gemini-pro' VERTEX_BASE_MODEL_ID = 'gemini-pro' PROXY_URL = 'https://test.apigee.net' @pytest.fixture def llm_request(): """Provides a sample LlmRequest for testing.""" return LlmRequest( model=APIGEE_GEMINI_MODEL_ID, contents=[ types.Content( role='user', parts=[types.Part.from_text(text='Test prompt')] ) ], ) @pytest.mark.asyncio @mock.patch('google.genai.Client') async def test_generate_content_async_non_streaming( mock_client_constructor, llm_request ): """Tests the generate_content_async method for non-streaming responses.""" apigee_llm_instance = ApigeeLlm( model=APIGEE_GEMINI_MODEL_ID, proxy_url=PROXY_URL, ) mock_client_instance = mock.Mock() mock_response = types.GenerateContentResponse( candidates=[ types.Candidate( content=Content( parts=[Part.from_text(text='Test response')], role='model', ) ) ] ) mock_client_instance.aio.models.generate_content = AsyncMock( return_value=mock_response ) mock_client_constructor.return_value = mock_client_instance response_generator = apigee_llm_instance.generate_content_async(llm_request) responses = [resp async for resp in response_generator] assert len(responses) == 1 llm_response = responses[0] assert llm_response.content.parts[0].text == 'Test response' assert llm_response.content.role == 'model' mock_client_constructor.assert_called_once() _, kwargs = mock_client_constructor.call_args assert not kwargs['enterprise'] http_options = kwargs['http_options'] assert http_options.base_url == PROXY_URL assert http_options.api_version == 'v1' assert 'user-agent' in http_options.headers assert 'x-goog-api-client' in http_options.headers mock_client_instance.aio.models.generate_content.assert_called_once_with( model=BASE_MODEL_ID, contents=llm_request.contents, config=llm_request.config, ) @pytest.mark.asyncio @mock.patch('google.genai.Client') async def test_generate_content_async_streaming( mock_client_constructor, llm_request ): """Tests the generate_content_async method for streaming responses.""" apigee_llm_instance = ApigeeLlm( model=APIGEE_GEMINI_MODEL_ID, proxy_url=PROXY_URL, ) mock_client_instance = mock.Mock() mock_responses = [ types.GenerateContentResponse( candidates=[ types.Candidate( content=Content( parts=[Part.from_text(text='Hello')], ) ) ] ), types.GenerateContentResponse( candidates=[ types.Candidate( content=Content( parts=[Part.from_text(text=',')], ) ) ] ), types.GenerateContentResponse( candidates=[ types.Candidate( content=Content( parts=[Part.from_text(text=' world!')], ) ) ] ), ] async def mock_stream_generator(): for r in mock_responses: yield r mock_client_instance.aio.models.generate_content_stream = AsyncMock( return_value=mock_stream_generator() ) mock_client_constructor.return_value = mock_client_instance response_generator = apigee_llm_instance.generate_content_async( llm_request, stream=True ) responses = [resp async for resp in response_generator] assert responses full_text_parts = [] for r in responses: for p in r.content.parts: if p.text: full_text_parts.append(p.text) full_text = ''.join(full_text_parts) assert 'Hello, world!' in full_text mock_client_instance.aio.models.generate_content_stream.assert_called_once_with( model=BASE_MODEL_ID, contents=llm_request.contents, config=llm_request.config, ) @pytest.mark.asyncio @mock.patch('google.genai.Client') async def test_generate_content_async_with_custom_headers( mock_client_constructor, llm_request ): """Tests that custom headers are passed in the request.""" custom_headers = { 'X-Custom-Header': 'custom-value', } apigee_llm = ApigeeLlm( model=APIGEE_GEMINI_MODEL_ID, proxy_url=PROXY_URL, custom_headers=custom_headers, ) mock_client_instance = mock.Mock() mock_response = types.GenerateContentResponse( candidates=[ types.Candidate( content=Content( parts=[Part.from_text(text='Test response')], role='model', ) ) ] ) mock_client_instance.aio.models.generate_content = AsyncMock( return_value=mock_response ) mock_client_constructor.return_value = mock_client_instance response_generator = apigee_llm.generate_content_async(llm_request) _ = [resp async for resp in response_generator] # Consume generator mock_client_constructor.assert_called_once() _, kwargs = mock_client_constructor.call_args http_options = kwargs['http_options'] assert http_options.headers['X-Custom-Header'] == 'custom-value' assert 'user-agent' in http_options.headers @pytest.mark.asyncio @mock.patch('google.genai.Client') async def test_vertex_model_path_parsing(mock_client_constructor): """Tests that Vertex AI model paths are parsed correctly.""" apigee_llm = ApigeeLlm(model=APIGEE_VERTEX_MODEL_ID, proxy_url=PROXY_URL) llm_request = LlmRequest( model=APIGEE_VERTEX_MODEL_ID, contents=[ types.Content( role='user', parts=[types.Part.from_text(text='Test prompt')] ) ], ) mock_client_instance = mock.Mock() mock_client_instance.aio.models.generate_content = AsyncMock( return_value=types.GenerateContentResponse( candidates=[ types.Candidate( content=Content( parts=[Part.from_text(text='Test response')], role='model', ) ) ] ) ) mock_client_constructor.return_value = mock_client_instance _ = [resp async for resp in apigee_llm.generate_content_async(llm_request)] mock_client_constructor.assert_called_once() _, kwargs = mock_client_constructor.call_args assert kwargs['enterprise'] assert kwargs['http_options'].api_version == 'v1beta' mock_client_instance.aio.models.generate_content.assert_called_once() call_kwargs = ( mock_client_instance.aio.models.generate_content.call_args.kwargs ) assert call_kwargs['model'] == VERTEX_BASE_MODEL_ID @pytest.mark.asyncio @mock.patch('google.genai.Client') async def test_proxy_url_from_env_variable(mock_client_constructor): """Tests that proxy_url is read from environment variable.""" with mock.patch.dict( os.environ, {'APIGEE_PROXY_URL': 'https://env.proxy.url'} ): apigee_llm = ApigeeLlm(model=APIGEE_GEMINI_MODEL_ID) llm_request = LlmRequest( model=APIGEE_GEMINI_MODEL_ID, contents=[ types.Content( role='user', parts=[types.Part.from_text(text='Test prompt')] ) ], ) mock_client_instance = mock.Mock() mock_client_instance.aio.models.generate_content = AsyncMock( return_value=types.GenerateContentResponse( candidates=[ types.Candidate( content=Content( parts=[Part.from_text(text='Test response')], role='model', ) ) ] ) ) mock_client_constructor.return_value = mock_client_instance _ = [resp async for resp in apigee_llm.generate_content_async(llm_request)] mock_client_constructor.assert_called_once() _, kwargs = mock_client_constructor.call_args assert kwargs['http_options'].base_url == 'https://env.proxy.url' @pytest.mark.parametrize( ('model_string', 'env_vars'), [ ( 'apigee/vertex_ai/gemini-2.5-flash', {'GOOGLE_CLOUD_LOCATION': 'test-location'}, ), ( 'apigee/vertex_ai/gemini-2.5-flash', {'GOOGLE_CLOUD_PROJECT': 'test-project'}, ), ( 'apigee/gemini-2.5-flash', { 'GOOGLE_GENAI_USE_ENTERPRISE': 'true', 'GOOGLE_CLOUD_LOCATION': 'test-location', }, ), ( 'apigee/gemini-2.5-flash', { 'GOOGLE_GENAI_USE_ENTERPRISE': 'true', 'GOOGLE_CLOUD_PROJECT': 'test-project', }, ), ], ) def test_vertex_model_missing_project_or_location_raises_error( model_string, env_vars ): """Tests that ValueError is raised for Vertex models if project or location is missing.""" with mock.patch.dict(os.environ, env_vars, clear=True): with pytest.raises(ValueError, match='environment variable must be set'): ApigeeLlm(model=model_string, proxy_url=PROXY_URL) @pytest.mark.asyncio @pytest.mark.parametrize( ( 'model_string', 'use_vertexai_env', 'expected_is_vertexai', 'expected_api_version', 'expected_model_id', ), [ ('apigee/gemini-2.5-flash', None, False, None, 'gemini-2.5-flash'), ('apigee/gemini-2.5-flash', 'true', True, None, 'gemini-2.5-flash'), ('apigee/gemini-2.5-flash', '1', True, None, 'gemini-2.5-flash'), ('apigee/gemini-2.5-flash', 'false', False, None, 'gemini-2.5-flash'), ('apigee/gemini-2.5-flash', '0', False, None, 'gemini-2.5-flash'), ( 'apigee/v1/gemini-2.5-flash', None, False, 'v1', 'gemini-2.5-flash', ), ( 'apigee/v1/gemini-2.5-flash', 'true', True, 'v1', 'gemini-2.5-flash', ), ( 'apigee/vertex_ai/gemini-2.5-flash', None, True, None, 'gemini-2.5-flash', ), ( 'apigee/vertex_ai/gemini-2.5-flash', 'false', True, None, 'gemini-2.5-flash', ), ( 'apigee/gemini/v1/gemini-2.5-flash', 'true', False, 'v1', 'gemini-2.5-flash', ), ( 'apigee/vertex_ai/v1beta/gemini-2.5-flash', 'false', True, 'v1beta', 'gemini-2.5-flash', ), ], ) @mock.patch('google.genai.Client') async def test_model_string_parsing_and_client_initialization( mock_client_constructor, model_string, use_vertexai_env, expected_is_vertexai, expected_api_version, expected_model_id, ): """Tests model string parsing and genai.Client initialization.""" env_vars = {} if use_vertexai_env is not None: env_vars['GOOGLE_GENAI_USE_ENTERPRISE'] = use_vertexai_env if expected_is_vertexai: env_vars['GOOGLE_CLOUD_PROJECT'] = 'test-project' env_vars['GOOGLE_CLOUD_LOCATION'] = 'test-location' # The ApigeeLlm is initialized in the 'with' block to make sure that the mock # of the environment variable is active. with mock.patch.dict(os.environ, env_vars, clear=True): apigee_llm = ApigeeLlm(model=model_string, proxy_url=PROXY_URL) request = LlmRequest(model=model_string, contents=[]) mock_client_instance = mock.Mock() mock_client_instance.aio.models.generate_content = AsyncMock( return_value=types.GenerateContentResponse( candidates=[ types.Candidate( content=Content(parts=[Part.from_text(text='')]) ) ] ) ) mock_client_constructor.return_value = mock_client_instance _ = [resp async for resp in apigee_llm.generate_content_async(request)] mock_client_constructor.assert_called_once() _, kwargs = mock_client_constructor.call_args assert kwargs['enterprise'] == expected_is_vertexai if expected_is_vertexai: assert kwargs['project'] == 'test-project' assert kwargs['location'] == 'test-location' http_options = kwargs['http_options'] assert http_options.api_version == expected_api_version ( mock_client_instance.aio.models.generate_content.assert_called_once_with( model=expected_model_id, contents=request.contents, config=request.config, ) ) @pytest.mark.asyncio @pytest.mark.parametrize( 'invalid_model_string', [ 'apigee/', # Missing model_id 'apigee', # Invalid format 'gemini-pro', # Invalid format 'apigee/vertex_ai/v1/model/extra', # Too many components 'apigee/unknown/model', ], ) async def test_invalid_model_strings_raise_value_error(invalid_model_string): """Tests that invalid model strings raise a ValueError.""" with pytest.raises( ValueError, match=f'Invalid model string: {invalid_model_string}' ): ApigeeLlm(model=invalid_model_string, proxy_url=PROXY_URL) @pytest.mark.asyncio @pytest.mark.parametrize( 'model', [ 'apigee/openai/gpt-4o', 'apigee/openai/v1/gpt-4o', 'apigee/openai/v1/gpt-3.5-turbo', ], ) async def test_validate_model_for_chat_completion_providers(model): """Tests that new providers like OpenAI are accepted.""" # Should not raise ValueError ApigeeLlm(model=model, proxy_url=PROXY_URL) @pytest.mark.parametrize( ('model', 'api_type', 'expected_api_type'), [ # Default case (input defaults to UNKNOWN) ( 'apigee/openai/gpt-4o', ApigeeLlm.ApiType.UNKNOWN, ApigeeLlm.ApiType.CHAT_COMPLETIONS, ), ( 'apigee/openai/v1/gpt-3.5-turbo', ApigeeLlm.ApiType.UNKNOWN, ApigeeLlm.ApiType.CHAT_COMPLETIONS, ), ( 'apigee/gemini/v1/gemini-pro', ApigeeLlm.ApiType.UNKNOWN, ApigeeLlm.ApiType.GENAI, ), ( 'apigee/vertex_ai/gemini-pro', ApigeeLlm.ApiType.UNKNOWN, ApigeeLlm.ApiType.GENAI, ), ( 'apigee/vertex_ai/v1beta/gemini-1.5-pro', ApigeeLlm.ApiType.UNKNOWN, ApigeeLlm.ApiType.GENAI, ), # Override by setting the ApiType ( 'apigee/gemini/pro', ApigeeLlm.ApiType.CHAT_COMPLETIONS, ApigeeLlm.ApiType.CHAT_COMPLETIONS, ), ( 'apigee/gemini/pro', ApigeeLlm.ApiType.GENAI, ApigeeLlm.ApiType.GENAI, ), ( 'apigee/openai/gpt-4o', ApigeeLlm.ApiType.CHAT_COMPLETIONS, ApigeeLlm.ApiType.CHAT_COMPLETIONS, ), ( 'apigee/openai/gpt-4o', ApigeeLlm.ApiType.GENAI, ApigeeLlm.ApiType.GENAI, ), # Override by setting the ApiType as a string ( 'apigee/gemini/pro', 'chat_completions', ApigeeLlm.ApiType.CHAT_COMPLETIONS, ), ( 'apigee/gemini/pro', 'genai', ApigeeLlm.ApiType.GENAI, ), ( 'apigee/openai/gpt-4o', 'chat_completions', ApigeeLlm.ApiType.CHAT_COMPLETIONS, ), ( 'apigee/openai/gpt-4o', 'genai', ApigeeLlm.ApiType.GENAI, ), ], ) def test_api_type_resolution(model, api_type, expected_api_type): """Tests that api_type is resolved correctly.""" llm = ApigeeLlm( model=model, proxy_url=PROXY_URL, api_type=api_type, ) assert llm._api_type == expected_api_type @pytest.mark.parametrize( ('input_value', 'expected_type'), [ ('chat_completions', ApigeeLlm.ApiType.CHAT_COMPLETIONS), ('genai', ApigeeLlm.ApiType.GENAI), ('unknown', ApigeeLlm.ApiType.UNKNOWN), ('', ApigeeLlm.ApiType.UNKNOWN), (None, ApigeeLlm.ApiType.UNKNOWN), ], ) def test_apitype_creation(input_value, expected_type): """Tests the creation of ApiType enum members.""" assert ApigeeLlm.ApiType(input_value) == expected_type def test_apitype_creation_invalid(): """Tests that invalid ApiType raises ValueError.""" with pytest.raises(ValueError): ApigeeLlm.ApiType('invalid') def test_invalid_api_type_raises_error(): """Tests that invalid string for api_type raises ValueError.""" with pytest.raises(ValueError): ApigeeLlm( model='apigee/gemini-pro', proxy_url=PROXY_URL, api_type='invalid_type', ) @pytest.mark.asyncio async def test_generate_content_async_dispatch_to_completions_client( llm_request, ): """Tests that generate_content_async uses CompletionsHTTPClient for OpenAI models.""" llm_request.model = 'apigee/openai/gpt-4o' with ( mock.patch.object( CompletionsHTTPClient, 'generate_content_async', ) as mock_completions_generate_content, mock.patch('google.genai.Client') as mock_genai_client, ): apigee_llm = ApigeeLlm(model='apigee/openai/gpt-4o', proxy_url=PROXY_URL) _ = [ r async for r in apigee_llm.generate_content_async( llm_request, stream=False ) ] mock_completions_generate_content.assert_called_once() mock_genai_client.assert_not_called() @pytest.mark.asyncio @pytest.mark.parametrize( 'model', [ 'apigee/openai/gpt-4o', 'apigee/openai/v1/gpt-3.5-turbo', ], ) async def test_api_key_injection_openai(model): """Tests that api_key is injected for OpenAI models.""" apigee_llm = ApigeeLlm( model=model, proxy_url=PROXY_URL, custom_headers={'Authorization': 'Bearer sk-test-key'}, ) client = apigee_llm._completions_http_client assert client._headers['Authorization'] == 'Bearer sk-test-key' def test_parse_response_usage_metadata(): """Tests that CompletionsHTTPClient parses usage metadata correctly including reasoning tokens.""" client = CompletionsHTTPClient(base_url='http://test') response_dict = { 'choices': [{ 'message': {'role': 'assistant', 'content': 'hello'}, 'finish_reason': 'stop', }], 'usage': { 'prompt_tokens': 10, 'completion_tokens': 5, 'total_tokens': 15, 'completion_tokens_details': {'reasoning_tokens': 4}, }, } llm_response = client._parse_response(response_dict) assert llm_response.usage_metadata.prompt_token_count == 10 assert llm_response.usage_metadata.candidates_token_count == 5 assert llm_response.usage_metadata.total_token_count == 15 assert llm_response.usage_metadata.thoughts_token_count == 4 @pytest.mark.asyncio @mock.patch('google.genai.Client') async def test_api_client_passes_credentials_when_provided( mock_client_constructor, llm_request ): """Tests that credentials passed to __init__ are forwarded to genai.Client.""" mock_credentials = mock.Mock() mock_client_instance = mock.Mock() mock_client_instance.aio.models.generate_content = AsyncMock( return_value=types.GenerateContentResponse( candidates=[ types.Candidate( content=Content( parts=[Part.from_text(text='Test response')], role='model', ) ) ] ) ) mock_client_constructor.return_value = mock_client_instance apigee_llm = ApigeeLlm( model=APIGEE_GEMINI_MODEL_ID, proxy_url=PROXY_URL, credentials=mock_credentials, ) _ = [resp async for resp in apigee_llm.generate_content_async(llm_request)] _, kwargs = mock_client_constructor.call_args assert kwargs['credentials'] is mock_credentials @pytest.mark.asyncio @mock.patch('google.genai.Client') async def test_api_client_omits_credentials_when_not_provided( mock_client_constructor, llm_request ): """Tests that credentials kwarg is not forwarded when not supplied.""" mock_client_instance = mock.Mock() mock_client_instance.aio.models.generate_content = AsyncMock( return_value=types.GenerateContentResponse( candidates=[ types.Candidate( content=Content( parts=[Part.from_text(text='Test response')], role='model', ) ) ] ) ) mock_client_constructor.return_value = mock_client_instance apigee_llm = ApigeeLlm( model=APIGEE_GEMINI_MODEL_ID, proxy_url=PROXY_URL, ) _ = [resp async for resp in apigee_llm.generate_content_async(llm_request)] _, kwargs = mock_client_constructor.call_args assert 'credentials' not in kwargs def test_parse_response_with_refusal(): """Tests that CompletionsHTTPClient parses refusal correctly.""" client = CompletionsHTTPClient(base_url='http://test') response_dict = { 'choices': [{ 'message': { 'role': 'assistant', 'refusal': 'I refuse to answer', }, 'finish_reason': 'stop', }], } llm_response = client._parse_response(response_dict) assert len(llm_response.content.parts) == 1 assert llm_response.content.parts[0].text == '[[REFUSAL]]: I refuse to answer' response_dict_mixed = { 'choices': [{ 'message': { 'role': 'assistant', 'content': 'Here is some content', 'refusal': 'But I refuse to answer the rest', }, 'finish_reason': 'stop', }], } llm_response_mixed = client._parse_response(response_dict_mixed) assert len(llm_response_mixed.content.parts) == 1 assert ( llm_response_mixed.content.parts[0].text == 'Here is some content\n[[REFUSAL]]: But I refuse to answer the rest' ) @pytest.mark.parametrize( ('parts', 'expected_message'), [ ( [ types.Part.from_text(text='[[REFUSAL]]: I refuse to answer'), types.Part.from_text(text='normal content'), ], { 'role': 'assistant', 'refusal': 'I refuse to answer', 'content': 'normal content', }, ), ( [ types.Part.from_text( text=( 'Here is some content\n[[REFUSAL]]: But I refuse to' ' answer the rest' ) ), ], { 'role': 'assistant', 'refusal': 'But I refuse to answer the rest', 'content': 'Here is some content', }, ), ], ) def test_construct_payload_with_refusal(parts, expected_message): """Tests that CompletionsHTTPClient constructs payload with refusal correctly.""" client = CompletionsHTTPClient(base_url='http://test') req = LlmRequest( model='apigee/openai/gpt-4o', contents=[ types.Content( role='model', parts=parts, ) ], ) payload = client._construct_payload(req, stream=False) messages = payload['messages'] assert messages == [expected_message]