From 6992d169219228b8a65edb820384e6d70a5e6859 Mon Sep 17 00:00:00 2001 From: Amy Wu Date: Fri, 17 Jul 2026 11:50:23 -0700 Subject: [PATCH] chore: internal testing PiperOrigin-RevId: 949684144 --- google/genai/_extra_utils.py | 32 +++-- .../genai/tests/afc/test_get_usage_header.py | 122 ++++++++++++++++++ 2 files changed, 142 insertions(+), 12 deletions(-) create mode 100644 google/genai/tests/afc/test_get_usage_header.py diff --git a/google/genai/_extra_utils.py b/google/genai/_extra_utils.py index 29191da11..20ad6a1fe 100644 --- a/google/genai/_extra_utils.py +++ b/google/genai/_extra_utils.py @@ -31,6 +31,7 @@ from . import _transformers as t from . import errors from . import types +from . import version as public_version from ._adapters import McpToGenAiToolAdapter @@ -701,7 +702,7 @@ def get_usage_header( usage: str = 'afc', ) -> types.GenerateContentConfig: """Sets the afc version label.""" - usage_header = f'google-genai-sdk/{usage}' + usage_header = f'google-genai-sdk/{public_version.__version__}+{usage}' if not config: config_model = types.GenerateContentConfig() elif isinstance(config, dict): @@ -712,15 +713,22 @@ def get_usage_header( if not config_model.http_options: config_model.http_options = types.HttpOptions() existing_headers = config_model.http_options.headers or {} - if 'user-agent' in existing_headers: - if usage_header not in existing_headers['user-agent']: - existing_headers['user-agent'] += f' {usage_header}' - else: - existing_headers['user-agent'] = usage_header - if 'x-goog-api-client' in existing_headers: - if usage_header not in existing_headers['x-goog-api-client']: - existing_headers['x-goog-api-client'] += f' {usage_header}' - else: - existing_headers['x-goog-api-client'] = usage_header + for header_key in ('user-agent', 'x-goog-api-client'): + if header_key in existing_headers: + if ( + f'+{usage}' not in existing_headers[header_key] + and usage_header not in existing_headers[header_key] + ): + if ( + f'google-genai-sdk/{public_version.__version__}' + in existing_headers[header_key] + ): + existing_headers[header_key] = existing_headers[header_key].replace( + f'google-genai-sdk/{public_version.__version__}', usage_header + ) + else: + existing_headers[header_key] += f' {usage_header}' + else: + existing_headers[header_key] = usage_header config_model.http_options.headers = existing_headers - return config_model \ No newline at end of file + return config_model diff --git a/google/genai/tests/afc/test_get_usage_header.py b/google/genai/tests/afc/test_get_usage_header.py new file mode 100644 index 000000000..e53ab9965 --- /dev/null +++ b/google/genai/tests/afc/test_get_usage_header.py @@ -0,0 +1,122 @@ +# 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. + +"""Unit tests for _extra_utils.get_usage_header.""" + +from ... import types +from ... import version as public_version +from ..._extra_utils import get_usage_header + + +def test_get_usage_header_none_config(): + config = get_usage_header(None) + expected_header = f'google-genai-sdk/{public_version.__version__}+afc' + assert config.http_options.headers['user-agent'] == expected_header + assert config.http_options.headers['x-goog-api-client'] == expected_header + + +def test_get_usage_header_dict_config(): + config = get_usage_header({'temperature': 0.5}, usage='afc') + expected_header = f'google-genai-sdk/{public_version.__version__}+afc' + assert config.temperature == 0.5 + assert config.http_options.headers['user-agent'] == expected_header + assert config.http_options.headers['x-goog-api-client'] == expected_header + + +def test_get_usage_header_custom_usage(): + config = get_usage_header(None, usage='chat') + expected_header = f'google-genai-sdk/{public_version.__version__}+chat' + assert config.http_options.headers['user-agent'] == expected_header + assert config.http_options.headers['x-goog-api-client'] == expected_header + + +def test_get_usage_header_with_existing_sdk_headers(): + initial_headers = { + 'user-agent': ( + f'google-genai-sdk/{public_version.__version__} gl-python/3.12' + ), + 'x-goog-api-client': ( + f'google-genai-sdk/{public_version.__version__} gl-python/3.12' + ), + } + config = types.GenerateContentConfig( + http_options=types.HttpOptions(headers=initial_headers) + ) + config = get_usage_header(config, usage='afc') + expected_header = ( + f'google-genai-sdk/{public_version.__version__}+afc gl-python/3.12' + ) + assert config.http_options.headers['user-agent'] == expected_header + assert config.http_options.headers['x-goog-api-client'] == expected_header + + +def test_get_usage_header_idempotent_no_duplicate_usage(): + config = get_usage_header(None, usage='afc') + config = get_usage_header(config, usage='afc') + expected_header = f'google-genai-sdk/{public_version.__version__}+afc' + assert config.http_options.headers['user-agent'] == expected_header + assert config.http_options.headers['x-goog-api-client'] == expected_header + + +def test_get_usage_header_multiple_usages_no_duplicate(): + config = get_usage_header(None, usage='chat') + config = get_usage_header(config, usage='afc') + # Call again with chat and afc to ensure no duplicate entries are appended + config = get_usage_header(config, usage='chat') + config = get_usage_header(config, usage='afc') + + user_agent = config.http_options.headers['user-agent'] + x_goog_api_client = config.http_options.headers['x-goog-api-client'] + + assert '+chat' in user_agent + assert '+afc' in user_agent + assert user_agent.count('+chat') == 1 + assert user_agent.count('+afc') == 1 + + assert '+chat' in x_goog_api_client + assert '+afc' in x_goog_api_client + assert x_goog_api_client.count('+chat') == 1 + assert x_goog_api_client.count('+afc') == 1 + + +def test_get_usage_header_with_custom_user_agent(): + initial_headers = { + 'user-agent': 'custom-agent/1.0', + 'x-goog-api-client': 'custom-agent/1.0', + } + config = types.GenerateContentConfig( + http_options=types.HttpOptions(headers=initial_headers) + ) + config = get_usage_header(config, usage='afc') + # Call again to ensure idempotency + config = get_usage_header(config, usage='afc') + + expected_usage = f'google-genai-sdk/{public_version.__version__}+afc' + assert 'custom-agent/1.0' in config.http_options.headers['user-agent'] + assert expected_usage in config.http_options.headers['user-agent'] + assert config.http_options.headers['user-agent'].count('+afc') == 1 + + +def test_get_usage_header_preserves_other_http_options(): + config = types.GenerateContentConfig( + http_options=types.HttpOptions( + api_version='v1alpha', + headers={'custom-header': 'value'}, + ) + ) + config = get_usage_header(config, usage='afc') + assert config.http_options.api_version == 'v1alpha' + assert config.http_options.headers['custom-header'] == 'value' + assert '+afc' in config.http_options.headers['user-agent'] + assert '+afc' in config.http_options.headers['x-goog-api-client']