mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
204 lines
7.4 KiB
Python
204 lines
7.4 KiB
Python
# 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.
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from unittest import mock
|
|
|
|
import google.adk
|
|
from google.adk.tools.bigquery.client import get_bigquery_client
|
|
import google.auth
|
|
from google.auth.exceptions import DefaultCredentialsError
|
|
from google.cloud.bigquery import client as bigquery_client
|
|
from google.oauth2.credentials import Credentials
|
|
|
|
|
|
def test_bigquery_client_default():
|
|
"""Test the default BigQuery client properties."""
|
|
# Trigger the BigQuery client creation
|
|
client = get_bigquery_client(
|
|
project="test-gcp-project",
|
|
credentials=mock.create_autospec(Credentials, instance=True),
|
|
)
|
|
|
|
# Verify that the client has the desired project set
|
|
assert client.project == "test-gcp-project"
|
|
assert client.location is None
|
|
|
|
|
|
def test_bigquery_client_project_set_explicit():
|
|
"""Test BigQuery client creation does not invoke default auth."""
|
|
# Let's simulate that no environment variables are set, so that any project
|
|
# set in there does not interfere with this test
|
|
with mock.patch.dict(os.environ, {}, clear=True):
|
|
with mock.patch.object(
|
|
google.auth, "default", autospec=True
|
|
) as mock_default_auth:
|
|
# Simulate exception from default auth
|
|
mock_default_auth.side_effect = DefaultCredentialsError(
|
|
"Your default credentials were not found"
|
|
)
|
|
|
|
# Trigger the BigQuery client creation
|
|
client = get_bigquery_client(
|
|
project="test-gcp-project",
|
|
credentials=mock.create_autospec(Credentials, instance=True),
|
|
)
|
|
|
|
# If we are here that already means client creation did not call default
|
|
# auth (otherwise we would have run into DefaultCredentialsError set
|
|
# above). For the sake of explicitness, trivially assert that the default
|
|
# auth was not called, and yet the project was set correctly
|
|
mock_default_auth.assert_not_called()
|
|
assert client.project == "test-gcp-project"
|
|
|
|
|
|
def test_bigquery_client_project_set_with_default_auth():
|
|
"""Test BigQuery client creation invokes default auth to set the project."""
|
|
# Let's simulate that no environment variables are set, so that any project
|
|
# set in there does not interfere with this test
|
|
with mock.patch.dict(os.environ, {}, clear=True):
|
|
with mock.patch.object(
|
|
google.auth, "default", autospec=True
|
|
) as mock_default_auth:
|
|
# Simulate credentials
|
|
mock_creds = mock.create_autospec(Credentials, instance=True)
|
|
|
|
# Simulate output of the default auth
|
|
mock_default_auth.return_value = (mock_creds, "test-gcp-project")
|
|
|
|
# Trigger the BigQuery client creation
|
|
client = get_bigquery_client(
|
|
project=None,
|
|
credentials=mock_creds,
|
|
)
|
|
|
|
# Verify that default auth was called once to set the client project
|
|
mock_default_auth.assert_called_once()
|
|
assert client.project == "test-gcp-project"
|
|
|
|
|
|
def test_bigquery_client_project_set_with_env():
|
|
"""Test BigQuery client creation sets the project from environment variable."""
|
|
# Let's simulate the project set in environment variables
|
|
with mock.patch.dict(
|
|
os.environ, {"GOOGLE_CLOUD_PROJECT": "test-gcp-project"}, clear=True
|
|
):
|
|
with mock.patch.object(
|
|
google.auth, "default", autospec=True
|
|
) as mock_default_auth:
|
|
# Simulate exception from default auth
|
|
mock_default_auth.side_effect = DefaultCredentialsError(
|
|
"Your default credentials were not found"
|
|
)
|
|
|
|
# Trigger the BigQuery client creation
|
|
client = get_bigquery_client(
|
|
project=None,
|
|
credentials=mock.create_autospec(Credentials, instance=True),
|
|
)
|
|
|
|
# If we are here that already means client creation did not call default
|
|
# auth (otherwise we would have run into DefaultCredentialsError set
|
|
# above). For the sake of explicitness, trivially assert that the default
|
|
# auth was not called, and yet the project was set correctly
|
|
mock_default_auth.assert_not_called()
|
|
assert client.project == "test-gcp-project"
|
|
|
|
|
|
def test_bigquery_client_user_agent_default():
|
|
"""Test BigQuery client default user agent."""
|
|
with mock.patch.object(
|
|
bigquery_client, "Connection", autospec=True
|
|
) as mock_connection:
|
|
# Trigger the BigQuery client creation
|
|
get_bigquery_client(
|
|
project="test-gcp-project",
|
|
credentials=mock.create_autospec(Credentials, instance=True),
|
|
)
|
|
|
|
# Verify that the tracking user agent was set
|
|
client_info_arg = mock_connection.call_args[1].get("client_info")
|
|
assert client_info_arg is not None
|
|
expected_user_agents = {
|
|
"adk-bigquery-tool",
|
|
f"google-adk/{google.adk.__version__}",
|
|
}
|
|
actual_user_agents = set(client_info_arg.user_agent.split())
|
|
assert expected_user_agents.issubset(actual_user_agents)
|
|
|
|
|
|
def test_bigquery_client_user_agent_custom():
|
|
"""Test BigQuery client custom user agent."""
|
|
with mock.patch.object(
|
|
bigquery_client, "Connection", autospec=True
|
|
) as mock_connection:
|
|
# Trigger the BigQuery client creation
|
|
get_bigquery_client(
|
|
project="test-gcp-project",
|
|
credentials=mock.create_autospec(Credentials, instance=True),
|
|
user_agent="custom_user_agent",
|
|
)
|
|
|
|
# Verify that the tracking user agent was set
|
|
client_info_arg = mock_connection.call_args[1].get("client_info")
|
|
assert client_info_arg is not None
|
|
expected_user_agents = {
|
|
"adk-bigquery-tool",
|
|
f"google-adk/{google.adk.__version__}",
|
|
"custom_user_agent",
|
|
}
|
|
actual_user_agents = set(client_info_arg.user_agent.split())
|
|
assert expected_user_agents.issubset(actual_user_agents)
|
|
|
|
|
|
def test_bigquery_client_user_agent_custom_list():
|
|
"""Test BigQuery client custom user agent."""
|
|
with mock.patch.object(
|
|
bigquery_client, "Connection", autospec=True
|
|
) as mock_connection:
|
|
# Trigger the BigQuery client creation
|
|
get_bigquery_client(
|
|
project="test-gcp-project",
|
|
credentials=mock.create_autospec(Credentials, instance=True),
|
|
user_agent=["custom_user_agent1", "custom_user_agent2"],
|
|
)
|
|
|
|
# Verify that the tracking user agents were set
|
|
client_info_arg = mock_connection.call_args[1].get("client_info")
|
|
assert client_info_arg is not None
|
|
expected_user_agents = {
|
|
"adk-bigquery-tool",
|
|
f"google-adk/{google.adk.__version__}",
|
|
"custom_user_agent1",
|
|
"custom_user_agent2",
|
|
}
|
|
actual_user_agents = set(client_info_arg.user_agent.split())
|
|
assert expected_user_agents.issubset(actual_user_agents)
|
|
|
|
|
|
def test_bigquery_client_location_custom():
|
|
"""Test BigQuery client custom location."""
|
|
# Trigger the BigQuery client creation
|
|
client = get_bigquery_client(
|
|
project="test-gcp-project",
|
|
credentials=mock.create_autospec(Credentials, instance=True),
|
|
location="us-central1",
|
|
)
|
|
|
|
# Verify that the client has the desired project set
|
|
assert client.project == "test-gcp-project"
|
|
assert client.location == "us-central1"
|