mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: add Spanner first-party toolset (breaking change to BigQueryTool, consolidating into generic GoogleTool)
Spanner toolset support basic operations to interact with Spanner table metadata and query results. Consolidate BigQueryTool into generic GoogleTool, so that BigQueryToolset and SpannerToolset can share. PiperOrigin-RevId: 794259782
This commit is contained in:
Google Team Member
committed by
Copybara-Service
Copybara-Service
parent
10e3dfab1a
commit
1fc8d20ae8
@@ -1,464 +0,0 @@
|
||||
# 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.
|
||||
|
||||
|
||||
import json
|
||||
from unittest.mock import Mock
|
||||
from unittest.mock import patch
|
||||
|
||||
from google.adk.auth.auth_tool import AuthConfig
|
||||
from google.adk.tools.bigquery.bigquery_credentials import BIGQUERY_TOKEN_CACHE_KEY
|
||||
from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig
|
||||
from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsManager
|
||||
from google.adk.tools.tool_context import ToolContext
|
||||
from google.auth.credentials import Credentials as AuthCredentials
|
||||
from google.auth.exceptions import RefreshError
|
||||
# Mock the Google OAuth and API dependencies
|
||||
from google.oauth2.credentials import Credentials as OAuthCredentials
|
||||
import pytest
|
||||
|
||||
|
||||
class TestBigQueryCredentialsManager:
|
||||
"""Test suite for BigQueryCredentialsManager OAuth flow handling.
|
||||
|
||||
This class tests the complex credential management logic including
|
||||
credential validation, refresh, OAuth flow orchestration, and the
|
||||
new token caching functionality through tool_context.state.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tool_context(self):
|
||||
"""Create a mock ToolContext for testing.
|
||||
|
||||
The ToolContext is the interface between tools and the broader
|
||||
agent framework, handling OAuth flows and state management.
|
||||
Now includes state dictionary for testing caching behavior.
|
||||
"""
|
||||
context = Mock(spec=ToolContext)
|
||||
context.get_auth_response = Mock(return_value=None)
|
||||
context.request_credential = Mock()
|
||||
context.state = {}
|
||||
return context
|
||||
|
||||
@pytest.fixture
|
||||
def credentials_config(self):
|
||||
"""Create a basic credentials configuration for testing."""
|
||||
return BigQueryCredentialsConfig(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
scopes=["https://www.googleapis.com/auth/calendar"],
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def manager(self, credentials_config):
|
||||
"""Create a credentials manager instance for testing."""
|
||||
return BigQueryCredentialsManager(credentials_config)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("credentials_class",),
|
||||
[
|
||||
pytest.param(OAuthCredentials, id="oauth"),
|
||||
pytest.param(AuthCredentials, id="auth"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_valid_credentials_with_valid_existing_creds(
|
||||
self, manager, mock_tool_context, credentials_class
|
||||
):
|
||||
"""Test that valid existing credentials are returned immediately.
|
||||
|
||||
When credentials are already valid, no refresh or OAuth flow
|
||||
should be needed. This is the optimal happy path scenario.
|
||||
"""
|
||||
# Create mock credentials that are already valid
|
||||
mock_creds = Mock(spec=credentials_class)
|
||||
mock_creds.valid = True
|
||||
manager.credentials_config.credentials = mock_creds
|
||||
|
||||
result = await manager.get_valid_credentials(mock_tool_context)
|
||||
|
||||
assert result == mock_creds
|
||||
# Verify no OAuth flow was triggered
|
||||
mock_tool_context.get_auth_response.assert_not_called()
|
||||
mock_tool_context.request_credential.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("valid",),
|
||||
[
|
||||
pytest.param(False, id="invalid"),
|
||||
pytest.param(True, id="valid"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_valid_credentials_with_existing_non_oauth_creds(
|
||||
self, manager, mock_tool_context, valid
|
||||
):
|
||||
"""Test that existing non-oauth credentials are returned immediately.
|
||||
|
||||
When credentials are of non-oauth type, no refresh or OAuth flow
|
||||
is triggered irrespective of whether it is valid or not.
|
||||
"""
|
||||
# Create mock credentials that are already valid
|
||||
mock_creds = Mock(spec=AuthCredentials)
|
||||
mock_creds.valid = valid
|
||||
manager.credentials_config.credentials = mock_creds
|
||||
|
||||
result = await manager.get_valid_credentials(mock_tool_context)
|
||||
|
||||
assert result == mock_creds
|
||||
# Verify no OAuth flow was triggered
|
||||
mock_tool_context.get_auth_response.assert_not_called()
|
||||
mock_tool_context.request_credential.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_credentials_from_cache_when_none_in_manager(
|
||||
self, manager, mock_tool_context
|
||||
):
|
||||
"""Test retrieving credentials from tool_context cache when manager has none.
|
||||
|
||||
This tests the new caching functionality where credentials can be
|
||||
retrieved from the tool context state when the manager instance
|
||||
doesn't have them loaded.
|
||||
"""
|
||||
# Manager starts with no credentials
|
||||
manager.credentials_config.credentials = None
|
||||
|
||||
# Create mock cached credentials JSON that would be stored in cache
|
||||
mock_cached_creds_json = json.dumps({
|
||||
"token": "cached_token",
|
||||
"refresh_token": "cached_refresh_token",
|
||||
"client_id": "test_client_id",
|
||||
"client_secret": "test_client_secret",
|
||||
})
|
||||
|
||||
# Set up the tool context state to contain cached credentials
|
||||
mock_tool_context.state[BIGQUERY_TOKEN_CACHE_KEY] = mock_cached_creds_json
|
||||
|
||||
# Mock the Credentials.from_authorized_user_info method
|
||||
with patch(
|
||||
"google.oauth2.credentials.Credentials.from_authorized_user_info"
|
||||
) as mock_from_json:
|
||||
mock_creds = Mock(spec=OAuthCredentials)
|
||||
mock_creds.valid = True
|
||||
mock_from_json.return_value = mock_creds
|
||||
|
||||
result = await manager.get_valid_credentials(mock_tool_context)
|
||||
|
||||
# Verify credentials were created from cached JSON
|
||||
mock_from_json.assert_called_once_with(
|
||||
json.loads(mock_cached_creds_json), manager.credentials_config.scopes
|
||||
)
|
||||
# Verify loaded credentials were not cached into manager
|
||||
assert manager.credentials_config.credentials is None
|
||||
# Verify valid cached credentials were returned
|
||||
assert result == mock_creds
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_credentials_in_manager_or_cache(
|
||||
self, manager, mock_tool_context
|
||||
):
|
||||
"""Test OAuth flow when no credentials exist in manager or cache.
|
||||
|
||||
This tests the scenario where both the manager and cache are empty,
|
||||
requiring a new OAuth flow to be initiated.
|
||||
"""
|
||||
# Manager starts with no credentials
|
||||
manager.credentials_config.credentials = None
|
||||
# Cache is also empty (state dict doesn't contain the key)
|
||||
|
||||
result = await manager.get_valid_credentials(mock_tool_context)
|
||||
|
||||
# Should trigger OAuth flow and return None (flow in progress)
|
||||
assert result is None
|
||||
mock_tool_context.request_credential.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("google.auth.transport.requests.Request")
|
||||
async def test_refresh_cached_credentials_success(
|
||||
self, mock_request_class, manager, mock_tool_context
|
||||
):
|
||||
"""Test successful refresh of expired credentials retrieved from cache.
|
||||
|
||||
This tests the interaction between caching and refresh functionality,
|
||||
ensuring that expired cached credentials can be refreshed properly.
|
||||
"""
|
||||
# Manager starts with no default credentials
|
||||
manager.credentials_config.credentials = None
|
||||
|
||||
# Create mock cached credentials JSON
|
||||
mock_cached_creds_json = json.dumps({
|
||||
"token": "expired_token",
|
||||
"refresh_token": "valid_refresh_token",
|
||||
"client_id": "test_client_id",
|
||||
"client_secret": "test_client_secret",
|
||||
})
|
||||
|
||||
mock_refreshed_creds_json = json.dumps({
|
||||
"token": "new_token",
|
||||
"refresh_token": "valid_refresh_token",
|
||||
"client_id": "test_client_id",
|
||||
"client_secret": "test_client_secret",
|
||||
})
|
||||
|
||||
# Set up the tool context state to contain cached credentials
|
||||
mock_tool_context.state[BIGQUERY_TOKEN_CACHE_KEY] = mock_cached_creds_json
|
||||
|
||||
# Create expired cached credentials with refresh token
|
||||
mock_cached_creds = Mock(spec=OAuthCredentials)
|
||||
mock_cached_creds.valid = False
|
||||
mock_cached_creds.expired = True
|
||||
mock_cached_creds.refresh_token = "valid_refresh_token"
|
||||
mock_cached_creds.to_json.return_value = mock_refreshed_creds_json
|
||||
|
||||
# Mock successful refresh
|
||||
def mock_refresh(request):
|
||||
mock_cached_creds.valid = True
|
||||
|
||||
mock_cached_creds.refresh = Mock(side_effect=mock_refresh)
|
||||
|
||||
# Mock the Credentials.from_authorized_user_info method
|
||||
with patch(
|
||||
"google.oauth2.credentials.Credentials.from_authorized_user_info"
|
||||
) as mock_from_json:
|
||||
mock_from_json.return_value = mock_cached_creds
|
||||
|
||||
result = await manager.get_valid_credentials(mock_tool_context)
|
||||
|
||||
# Verify credentials were created from cached JSON
|
||||
mock_from_json.assert_called_once_with(
|
||||
json.loads(mock_cached_creds_json), manager.credentials_config.scopes
|
||||
)
|
||||
# Verify refresh was attempted and succeeded
|
||||
mock_cached_creds.refresh.assert_called_once()
|
||||
# Verify refreshed credentials were not cached into manager
|
||||
assert manager.credentials_config.credentials is None
|
||||
# Verify refreshed credentials were cached
|
||||
assert (
|
||||
"new_token"
|
||||
== json.loads(mock_tool_context.state[BIGQUERY_TOKEN_CACHE_KEY])[
|
||||
"token"
|
||||
]
|
||||
)
|
||||
assert result == mock_cached_creds
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("google.auth.transport.requests.Request")
|
||||
async def test_get_valid_credentials_with_refresh_success(
|
||||
self, mock_request_class, manager, mock_tool_context
|
||||
):
|
||||
"""Test successful credential refresh when tokens are expired.
|
||||
|
||||
This tests the automatic token refresh capability that prevents
|
||||
users from having to re-authenticate for every expired token.
|
||||
"""
|
||||
# Create expired credentials with refresh token
|
||||
mock_creds = Mock(spec=OAuthCredentials)
|
||||
mock_creds.valid = False
|
||||
mock_creds.expired = True
|
||||
mock_creds.refresh_token = "refresh_token"
|
||||
|
||||
# Mock successful refresh
|
||||
def mock_refresh(request):
|
||||
mock_creds.valid = True
|
||||
|
||||
mock_creds.refresh = Mock(side_effect=mock_refresh)
|
||||
manager.credentials_config.credentials = mock_creds
|
||||
|
||||
result = await manager.get_valid_credentials(mock_tool_context)
|
||||
|
||||
assert result == mock_creds
|
||||
mock_creds.refresh.assert_called_once()
|
||||
# Verify credentials were cached after successful refresh
|
||||
assert manager.credentials_config.credentials == mock_creds
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("google.auth.transport.requests.Request")
|
||||
async def test_get_valid_credentials_with_refresh_failure(
|
||||
self, mock_request_class, manager, mock_tool_context
|
||||
):
|
||||
"""Test OAuth flow trigger when credential refresh fails.
|
||||
|
||||
When refresh tokens expire or become invalid, the system should
|
||||
gracefully fall back to requesting a new OAuth flow.
|
||||
"""
|
||||
# Create expired credentials that fail to refresh
|
||||
mock_creds = Mock(spec=OAuthCredentials)
|
||||
mock_creds.valid = False
|
||||
mock_creds.expired = True
|
||||
mock_creds.refresh_token = "expired_refresh_token"
|
||||
mock_creds.refresh = Mock(side_effect=RefreshError("Refresh failed"))
|
||||
manager.credentials_config.credentials = mock_creds
|
||||
|
||||
result = await manager.get_valid_credentials(mock_tool_context)
|
||||
|
||||
# Should trigger OAuth flow and return None (flow in progress)
|
||||
assert result is None
|
||||
mock_tool_context.request_credential.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oauth_flow_completion_with_caching(
|
||||
self, manager, mock_tool_context
|
||||
):
|
||||
"""Test successful OAuth flow completion with proper credential caching.
|
||||
|
||||
This tests the happy path where a user completes the OAuth flow
|
||||
and the system successfully creates and caches new credentials
|
||||
in both the manager and the tool context state.
|
||||
"""
|
||||
# Mock OAuth response indicating completed flow
|
||||
mock_auth_response = Mock()
|
||||
mock_auth_response.oauth2.access_token = "new_access_token"
|
||||
mock_auth_response.oauth2.refresh_token = "new_refresh_token"
|
||||
mock_tool_context.get_auth_response.return_value = mock_auth_response
|
||||
|
||||
# Create a mock credentials instance that will represent our created credentials
|
||||
mock_creds = Mock(spec=OAuthCredentials)
|
||||
# Make the JSON match what a real Credentials object would produce
|
||||
mock_creds_json = (
|
||||
'{"token": "new_access_token", "refresh_token": "new_refresh_token",'
|
||||
' "token_uri": "https://oauth2.googleapis.com/token", "client_id":'
|
||||
' "test_client_id", "client_secret": "test_client_secret", "scopes":'
|
||||
' ["https://www.googleapis.com/auth/calendar"], "universe_domain":'
|
||||
' "googleapis.com", "account": ""}'
|
||||
)
|
||||
mock_creds.to_json.return_value = mock_creds_json
|
||||
|
||||
# Use the full module path as it appears in the project structure
|
||||
with patch(
|
||||
"google.adk.tools.bigquery.bigquery_credentials.google.oauth2.credentials.Credentials",
|
||||
return_value=mock_creds,
|
||||
) as mock_credentials_class:
|
||||
result = await manager.get_valid_credentials(mock_tool_context)
|
||||
|
||||
# Verify new credentials were created
|
||||
assert result == mock_creds
|
||||
# Verify credentials are created with correct parameters
|
||||
mock_credentials_class.assert_called_once()
|
||||
call_kwargs = mock_credentials_class.call_args[1]
|
||||
assert call_kwargs["token"] == "new_access_token"
|
||||
assert call_kwargs["refresh_token"] == "new_refresh_token"
|
||||
|
||||
# Verify credentials are not cached in manager
|
||||
assert manager.credentials_config.credentials is None
|
||||
# Verify credentials are also cached in tool context state
|
||||
assert (
|
||||
mock_tool_context.state[BIGQUERY_TOKEN_CACHE_KEY] == mock_creds_json
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oauth_flow_in_progress(self, manager, mock_tool_context):
|
||||
"""Test OAuth flow initiation when no auth response is available.
|
||||
|
||||
This tests the case where the OAuth flow needs to be started,
|
||||
and the user hasn't completed authorization yet.
|
||||
"""
|
||||
# No existing credentials, no auth response (flow not completed)
|
||||
manager.credentials_config.credentials = None
|
||||
mock_tool_context.get_auth_response.return_value = None
|
||||
|
||||
result = await manager.get_valid_credentials(mock_tool_context)
|
||||
|
||||
# Should return None and request credential flow
|
||||
assert result is None
|
||||
mock_tool_context.request_credential.assert_called_once()
|
||||
|
||||
# Verify the auth configuration includes correct scopes and endpoints
|
||||
call_args = mock_tool_context.request_credential.call_args[0][0]
|
||||
assert isinstance(call_args, AuthConfig)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_persistence_across_manager_instances(
|
||||
self, credentials_config, mock_tool_context
|
||||
):
|
||||
"""Test that cached credentials persist across different manager instances.
|
||||
|
||||
This tests the key benefit of the tool context caching - that
|
||||
credentials can be shared between different instances of the
|
||||
credential manager, avoiding redundant OAuth flows.
|
||||
"""
|
||||
# Create first manager instance and simulate OAuth completion
|
||||
manager1 = BigQueryCredentialsManager(credentials_config)
|
||||
|
||||
# Mock OAuth response for first manager
|
||||
mock_auth_response = Mock()
|
||||
mock_auth_response.oauth2.access_token = "cached_access_token"
|
||||
mock_auth_response.oauth2.refresh_token = "cached_refresh_token"
|
||||
mock_tool_context.get_auth_response.return_value = mock_auth_response
|
||||
|
||||
# Create the mock credentials instance that will be returned by the constructor
|
||||
mock_creds = Mock(spec=OAuthCredentials)
|
||||
# Make sure our mock JSON matches the structure that real Credentials objects produce
|
||||
mock_creds_json = (
|
||||
'{"token": "cached_access_token", "refresh_token":'
|
||||
' "cached_refresh_token", "token_uri":'
|
||||
' "https://oauth2.googleapis.com/token", "client_id": "test_client_id",'
|
||||
' "client_secret": "test_client_secret", "scopes":'
|
||||
' ["https://www.googleapis.com/auth/calendar"], "universe_domain":'
|
||||
' "googleapis.com", "account": ""}'
|
||||
)
|
||||
mock_creds.to_json.return_value = mock_creds_json
|
||||
mock_creds.valid = True
|
||||
|
||||
# Use the correct module path - without the 'src.' prefix
|
||||
with patch(
|
||||
"google.adk.tools.bigquery.bigquery_credentials.google.oauth2.credentials.Credentials",
|
||||
return_value=mock_creds,
|
||||
) as mock_credentials_class:
|
||||
# Complete OAuth flow with first manager
|
||||
result1 = await manager1.get_valid_credentials(mock_tool_context)
|
||||
|
||||
# Verify credentials were cached in tool context
|
||||
assert BIGQUERY_TOKEN_CACHE_KEY in mock_tool_context.state
|
||||
cached_creds_json = mock_tool_context.state[BIGQUERY_TOKEN_CACHE_KEY]
|
||||
assert cached_creds_json == mock_creds_json
|
||||
|
||||
# Create second manager instance (simulating new request/session)
|
||||
manager2 = BigQueryCredentialsManager(credentials_config)
|
||||
credentials_config.credentials = None
|
||||
|
||||
# Reset auth response to None (no new OAuth flow available)
|
||||
mock_tool_context.get_auth_response.return_value = None
|
||||
|
||||
# Mock the from_authorized_user_info method for the second manager
|
||||
with patch(
|
||||
"google.adk.tools.bigquery.bigquery_credentials.google.oauth2.credentials.Credentials.from_authorized_user_info"
|
||||
) as mock_from_json:
|
||||
mock_cached_creds = Mock(spec=OAuthCredentials)
|
||||
mock_cached_creds.valid = True
|
||||
mock_from_json.return_value = mock_cached_creds
|
||||
|
||||
# Get credentials with second manager
|
||||
result2 = await manager2.get_valid_credentials(mock_tool_context)
|
||||
|
||||
# Verify second manager retrieved cached credentials successfully
|
||||
assert result2 == mock_cached_creds
|
||||
assert manager2.credentials_config.credentials is None
|
||||
assert (
|
||||
cached_creds_json == mock_tool_context.state[BIGQUERY_TOKEN_CACHE_KEY]
|
||||
)
|
||||
# The from_authorized_user_info should be called with the complete JSON structure
|
||||
mock_from_json.assert_called_once()
|
||||
# Extract the actual argument that was passed to verify it's the right JSON structure
|
||||
actual_json_arg = mock_from_json.call_args[0][0]
|
||||
# We need to parse and compare the structure rather than exact string match
|
||||
# since the order of keys in JSON might differ
|
||||
import json
|
||||
|
||||
expected_data = json.loads(mock_creds_json)
|
||||
actual_data = (
|
||||
actual_json_arg
|
||||
if isinstance(actual_json_arg, dict)
|
||||
else json.loads(actual_json_arg)
|
||||
)
|
||||
assert actual_data == expected_data
|
||||
@@ -74,8 +74,8 @@ def test_ask_data_insights_success(mock_get_stream):
|
||||
# 2. Create mock inputs for the function call
|
||||
mock_creds = mock.Mock()
|
||||
mock_creds.token = "fake-token"
|
||||
mock_config = mock.Mock()
|
||||
mock_config.max_query_result_rows = 100
|
||||
mock_settings = mock.Mock()
|
||||
mock_settings.max_query_result_rows = 100
|
||||
|
||||
# 3. Call the function under test
|
||||
result = data_insights_tool.ask_data_insights(
|
||||
@@ -83,7 +83,7 @@ def test_ask_data_insights_success(mock_get_stream):
|
||||
user_query_with_context="test query",
|
||||
table_references=[],
|
||||
credentials=mock_creds,
|
||||
config=mock_config,
|
||||
settings=mock_settings,
|
||||
)
|
||||
|
||||
# 4. Assert the results are as expected
|
||||
@@ -101,7 +101,7 @@ def test_ask_data_insights_handles_exception(mock_get_stream):
|
||||
# 2. Create mock inputs
|
||||
mock_creds = mock.Mock()
|
||||
mock_creds.token = "fake-token"
|
||||
mock_config = mock.Mock()
|
||||
mock_settings = mock.Mock()
|
||||
|
||||
# 3. Call the function
|
||||
result = data_insights_tool.ask_data_insights(
|
||||
@@ -109,7 +109,7 @@ def test_ask_data_insights_handles_exception(mock_get_stream):
|
||||
user_query_with_context="test query",
|
||||
table_references=[],
|
||||
credentials=mock_creds,
|
||||
config=mock_config,
|
||||
settings=mock_settings,
|
||||
)
|
||||
|
||||
# 4. Assert that the error was caught and formatted correctly
|
||||
|
||||
@@ -37,7 +37,7 @@ import pytest
|
||||
|
||||
|
||||
async def get_tool(
|
||||
name: str, tool_config: Optional[BigQueryToolConfig] = None
|
||||
name: str, tool_settings: Optional[BigQueryToolConfig] = None
|
||||
) -> BaseTool:
|
||||
"""Get a tool from BigQuery toolset.
|
||||
|
||||
@@ -54,7 +54,7 @@ async def get_tool(
|
||||
toolset = BigQueryToolset(
|
||||
credentials_config=credentials_config,
|
||||
tool_filter=[name],
|
||||
bigquery_tool_config=tool_config,
|
||||
bigquery_tool_config=tool_settings,
|
||||
)
|
||||
|
||||
tools = await toolset.get_tools()
|
||||
@@ -64,7 +64,7 @@ async def get_tool(
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tool_config",),
|
||||
("tool_settings",),
|
||||
[
|
||||
pytest.param(None, id="no-config"),
|
||||
pytest.param(BigQueryToolConfig(), id="default-config"),
|
||||
@@ -75,14 +75,14 @@ async def get_tool(
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sql_declaration_read_only(tool_config):
|
||||
async def test_execute_sql_declaration_read_only(tool_settings):
|
||||
"""Test BigQuery execute_sql tool declaration in read-only mode.
|
||||
|
||||
This test verifies that the execute_sql tool declaration reflects the
|
||||
read-only capability.
|
||||
"""
|
||||
tool_name = "execute_sql"
|
||||
tool = await get_tool(tool_name, tool_config)
|
||||
tool = await get_tool(tool_name, tool_settings)
|
||||
assert tool.name == tool_name
|
||||
assert tool.description == textwrap.dedent("""\
|
||||
Run a BigQuery or BigQuery ML SQL query in the project and return the result.
|
||||
@@ -92,7 +92,7 @@ async def test_execute_sql_declaration_read_only(tool_config):
|
||||
executed.
|
||||
query (str): The BigQuery SQL query to be executed.
|
||||
credentials (Credentials): The credentials to use for the request.
|
||||
config (BigQueryToolConfig): The configuration for the tool.
|
||||
settings (BigQueryToolConfig): The settings for the tool.
|
||||
tool_context (ToolContext): The context for the tool.
|
||||
|
||||
Returns:
|
||||
@@ -127,7 +127,7 @@ async def test_execute_sql_declaration_read_only(tool_config):
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tool_config",),
|
||||
("tool_settings",),
|
||||
[
|
||||
pytest.param(
|
||||
BigQueryToolConfig(write_mode=WriteMode.ALLOWED),
|
||||
@@ -136,14 +136,14 @@ async def test_execute_sql_declaration_read_only(tool_config):
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sql_declaration_write(tool_config):
|
||||
async def test_execute_sql_declaration_write(tool_settings):
|
||||
"""Test BigQuery execute_sql tool declaration with all writes enabled.
|
||||
|
||||
This test verifies that the execute_sql tool declaration reflects the write
|
||||
capability.
|
||||
"""
|
||||
tool_name = "execute_sql"
|
||||
tool = await get_tool(tool_name, tool_config)
|
||||
tool = await get_tool(tool_name, tool_settings)
|
||||
assert tool.name == tool_name
|
||||
assert tool.description == textwrap.dedent("""\
|
||||
Run a BigQuery or BigQuery ML SQL query in the project and return the result.
|
||||
@@ -153,7 +153,7 @@ async def test_execute_sql_declaration_write(tool_config):
|
||||
executed.
|
||||
query (str): The BigQuery SQL query to be executed.
|
||||
credentials (Credentials): The credentials to use for the request.
|
||||
config (BigQueryToolConfig): The configuration for the tool.
|
||||
settings (BigQueryToolConfig): The settings for the tool.
|
||||
tool_context (ToolContext): The context for the tool.
|
||||
|
||||
Returns:
|
||||
@@ -326,7 +326,7 @@ async def test_execute_sql_declaration_write(tool_config):
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tool_config",),
|
||||
("tool_settings",),
|
||||
[
|
||||
pytest.param(
|
||||
BigQueryToolConfig(write_mode=WriteMode.PROTECTED),
|
||||
@@ -335,14 +335,14 @@ async def test_execute_sql_declaration_write(tool_config):
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sql_declaration_protected_write(tool_config):
|
||||
async def test_execute_sql_declaration_protected_write(tool_settings):
|
||||
"""Test BigQuery execute_sql tool declaration with protected writes enabled.
|
||||
|
||||
This test verifies that the execute_sql tool declaration reflects the
|
||||
protected write capability.
|
||||
"""
|
||||
tool_name = "execute_sql"
|
||||
tool = await get_tool(tool_name, tool_config)
|
||||
tool = await get_tool(tool_name, tool_settings)
|
||||
assert tool.name == tool_name
|
||||
assert tool.description == textwrap.dedent("""\
|
||||
Run a BigQuery or BigQuery ML SQL query in the project and return the result.
|
||||
@@ -352,7 +352,7 @@ async def test_execute_sql_declaration_protected_write(tool_config):
|
||||
executed.
|
||||
query (str): The BigQuery SQL query to be executed.
|
||||
credentials (Credentials): The credentials to use for the request.
|
||||
config (BigQueryToolConfig): The configuration for the tool.
|
||||
settings (BigQueryToolConfig): The settings for the tool.
|
||||
tool_context (ToolContext): The context for the tool.
|
||||
|
||||
Returns:
|
||||
@@ -530,7 +530,7 @@ def test_execute_sql_select_stmt(write_mode):
|
||||
statement_type = "SELECT"
|
||||
query_result = [{"num": 123}]
|
||||
credentials = mock.create_autospec(Credentials, instance=True)
|
||||
tool_config = BigQueryToolConfig(write_mode=write_mode)
|
||||
tool_settings = BigQueryToolConfig(write_mode=write_mode)
|
||||
tool_context = mock.create_autospec(ToolContext, instance=True)
|
||||
tool_context.state.get.return_value = (
|
||||
"test-bq-session-id",
|
||||
@@ -550,7 +550,9 @@ def test_execute_sql_select_stmt(write_mode):
|
||||
bq_client.query_and_wait.return_value = query_result
|
||||
|
||||
# Test the tool
|
||||
result = execute_sql(project, query, credentials, tool_config, tool_context)
|
||||
result = execute_sql(
|
||||
project, query, credentials, tool_settings, tool_context
|
||||
)
|
||||
assert result == {"status": "SUCCESS", "rows": query_result}
|
||||
|
||||
|
||||
@@ -586,7 +588,7 @@ def test_execute_sql_non_select_stmt_write_allowed(query, statement_type):
|
||||
project = "my_project"
|
||||
query_result = []
|
||||
credentials = mock.create_autospec(Credentials, instance=True)
|
||||
tool_config = BigQueryToolConfig(write_mode=WriteMode.ALLOWED)
|
||||
tool_settings = BigQueryToolConfig(write_mode=WriteMode.ALLOWED)
|
||||
tool_context = mock.create_autospec(ToolContext, instance=True)
|
||||
|
||||
with mock.patch("google.cloud.bigquery.Client", autospec=False) as Client:
|
||||
@@ -602,7 +604,9 @@ def test_execute_sql_non_select_stmt_write_allowed(query, statement_type):
|
||||
bq_client.query_and_wait.return_value = query_result
|
||||
|
||||
# Test the tool
|
||||
result = execute_sql(project, query, credentials, tool_config, tool_context)
|
||||
result = execute_sql(
|
||||
project, query, credentials, tool_settings, tool_context
|
||||
)
|
||||
assert result == {"status": "SUCCESS", "rows": query_result}
|
||||
|
||||
|
||||
@@ -638,7 +642,7 @@ def test_execute_sql_non_select_stmt_write_blocked(query, statement_type):
|
||||
project = "my_project"
|
||||
query_result = []
|
||||
credentials = mock.create_autospec(Credentials, instance=True)
|
||||
tool_config = BigQueryToolConfig(write_mode=WriteMode.BLOCKED)
|
||||
tool_settings = BigQueryToolConfig(write_mode=WriteMode.BLOCKED)
|
||||
tool_context = mock.create_autospec(ToolContext, instance=True)
|
||||
|
||||
with mock.patch("google.cloud.bigquery.Client", autospec=False) as Client:
|
||||
@@ -654,7 +658,9 @@ def test_execute_sql_non_select_stmt_write_blocked(query, statement_type):
|
||||
bq_client.query_and_wait.return_value = query_result
|
||||
|
||||
# Test the tool
|
||||
result = execute_sql(project, query, credentials, tool_config, tool_context)
|
||||
result = execute_sql(
|
||||
project, query, credentials, tool_settings, tool_context
|
||||
)
|
||||
assert result == {
|
||||
"status": "ERROR",
|
||||
"error_details": "Read-only mode only supports SELECT statements.",
|
||||
@@ -693,7 +699,7 @@ def test_execute_sql_non_select_stmt_write_protected(query, statement_type):
|
||||
project = "my_project"
|
||||
query_result = []
|
||||
credentials = mock.create_autospec(Credentials, instance=True)
|
||||
tool_config = BigQueryToolConfig(write_mode=WriteMode.PROTECTED)
|
||||
tool_settings = BigQueryToolConfig(write_mode=WriteMode.PROTECTED)
|
||||
tool_context = mock.create_autospec(ToolContext, instance=True)
|
||||
tool_context.state.get.return_value = (
|
||||
"test-bq-session-id",
|
||||
@@ -714,7 +720,9 @@ def test_execute_sql_non_select_stmt_write_protected(query, statement_type):
|
||||
bq_client.query_and_wait.return_value = query_result
|
||||
|
||||
# Test the tool
|
||||
result = execute_sql(project, query, credentials, tool_config, tool_context)
|
||||
result = execute_sql(
|
||||
project, query, credentials, tool_settings, tool_context
|
||||
)
|
||||
assert result == {"status": "SUCCESS", "rows": query_result}
|
||||
|
||||
|
||||
@@ -756,7 +764,7 @@ def test_execute_sql_non_select_stmt_write_protected_persistent_target(
|
||||
project = "my_project"
|
||||
query_result = []
|
||||
credentials = mock.create_autospec(Credentials, instance=True)
|
||||
tool_config = BigQueryToolConfig(write_mode=WriteMode.PROTECTED)
|
||||
tool_settings = BigQueryToolConfig(write_mode=WriteMode.PROTECTED)
|
||||
tool_context = mock.create_autospec(ToolContext, instance=True)
|
||||
tool_context.state.get.return_value = (
|
||||
"test-bq-session-id",
|
||||
@@ -777,7 +785,9 @@ def test_execute_sql_non_select_stmt_write_protected_persistent_target(
|
||||
bq_client.query_and_wait.return_value = query_result
|
||||
|
||||
# Test the tool
|
||||
result = execute_sql(project, query, credentials, tool_config, tool_context)
|
||||
result = execute_sql(
|
||||
project, query, credentials, tool_settings, tool_context
|
||||
)
|
||||
assert result == {
|
||||
"status": "ERROR",
|
||||
"error_details": (
|
||||
@@ -808,7 +818,7 @@ def test_execute_sql_no_default_auth(
|
||||
statement_type = "SELECT"
|
||||
query_result = [{"num": 123}]
|
||||
credentials = mock.create_autospec(Credentials, instance=True)
|
||||
tool_config = BigQueryToolConfig(write_mode=write_mode)
|
||||
tool_settings = BigQueryToolConfig(write_mode=write_mode)
|
||||
tool_context = mock.create_autospec(ToolContext, instance=True)
|
||||
tool_context.state.get.return_value = (
|
||||
"test-bq-session-id",
|
||||
@@ -830,7 +840,7 @@ def test_execute_sql_no_default_auth(
|
||||
mock_query_and_wait.return_value = query_result
|
||||
|
||||
# Test the tool worked without invoking default auth
|
||||
result = execute_sql(project, query, credentials, tool_config, tool_context)
|
||||
result = execute_sql(project, query, credentials, tool_settings, tool_context)
|
||||
assert result == {"status": "SUCCESS", "rows": query_result}
|
||||
mock_default_auth.assert_not_called()
|
||||
|
||||
@@ -959,7 +969,7 @@ def test_execute_sql_result_dtype(
|
||||
project = "my_project"
|
||||
statement_type = "SELECT"
|
||||
credentials = mock.create_autospec(Credentials, instance=True)
|
||||
tool_config = BigQueryToolConfig()
|
||||
tool_settings = BigQueryToolConfig()
|
||||
tool_context = mock.create_autospec(ToolContext, instance=True)
|
||||
|
||||
# Simulate the result of query API
|
||||
@@ -971,5 +981,5 @@ def test_execute_sql_result_dtype(
|
||||
mock_query_and_wait.return_value = query_result
|
||||
|
||||
# Test the tool worked without invoking default auth
|
||||
result = execute_sql(project, query, credentials, tool_config, tool_context)
|
||||
result = execute_sql(project, query, credentials, tool_settings, tool_context)
|
||||
assert result == {"status": "SUCCESS", "rows": tool_result_rows}
|
||||
|
||||
@@ -1,302 +0,0 @@
|
||||
# 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 unittest.mock import Mock
|
||||
from unittest.mock import patch
|
||||
|
||||
from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig
|
||||
from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsManager
|
||||
from google.adk.tools.bigquery.bigquery_tool import BigQueryTool
|
||||
from google.adk.tools.bigquery.config import BigQueryToolConfig
|
||||
from google.adk.tools.tool_context import ToolContext
|
||||
# Mock the Google OAuth and API dependencies
|
||||
from google.oauth2.credentials import Credentials
|
||||
import pytest
|
||||
|
||||
|
||||
class TestBigQueryTool:
|
||||
"""Test suite for BigQueryTool OAuth integration and execution.
|
||||
|
||||
This class tests the high-level tool execution logic that combines
|
||||
credential management with actual function execution.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tool_context(self):
|
||||
"""Create a mock ToolContext for testing tool execution."""
|
||||
context = Mock(spec=ToolContext)
|
||||
context.get_auth_response = Mock(return_value=None)
|
||||
context.request_credential = Mock()
|
||||
return context
|
||||
|
||||
@pytest.fixture
|
||||
def sample_function(self):
|
||||
"""Create a sample function that accepts credentials for testing.
|
||||
|
||||
This simulates a real Google API tool function that needs
|
||||
authenticated credentials to perform its work.
|
||||
"""
|
||||
|
||||
def sample_func(param1: str, credentials: Credentials = None) -> dict:
|
||||
"""Sample function that uses Google API credentials."""
|
||||
if credentials:
|
||||
return {"result": f"Success with {param1}", "authenticated": True}
|
||||
else:
|
||||
return {"result": f"Success with {param1}", "authenticated": False}
|
||||
|
||||
return sample_func
|
||||
|
||||
@pytest.fixture
|
||||
def async_sample_function(self):
|
||||
"""Create an async sample function for testing async execution paths."""
|
||||
|
||||
async def async_sample_func(
|
||||
param1: str, credentials: Credentials = None
|
||||
) -> dict:
|
||||
"""Async sample function that uses Google API credentials."""
|
||||
if credentials:
|
||||
return {"result": f"Async success with {param1}", "authenticated": True}
|
||||
else:
|
||||
return {
|
||||
"result": f"Async success with {param1}",
|
||||
"authenticated": False,
|
||||
}
|
||||
|
||||
return async_sample_func
|
||||
|
||||
@pytest.fixture
|
||||
def credentials_config(self):
|
||||
"""Create credentials configuration for testing."""
|
||||
return BigQueryCredentialsConfig(
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
scopes=["https://www.googleapis.com/auth/bigquery"],
|
||||
)
|
||||
|
||||
def test_tool_initialization_with_credentials(
|
||||
self, sample_function, credentials_config
|
||||
):
|
||||
"""Test that BigQueryTool initializes correctly with credentials.
|
||||
|
||||
The tool should properly inherit from FunctionTool while adding
|
||||
Google API specific credential management capabilities.
|
||||
"""
|
||||
tool = BigQueryTool(
|
||||
func=sample_function, credentials_config=credentials_config
|
||||
)
|
||||
|
||||
assert tool.func == sample_function
|
||||
assert tool._credentials_manager is not None
|
||||
assert isinstance(tool._credentials_manager, BigQueryCredentialsManager)
|
||||
# Verify that 'credentials' parameter is ignored in function signature analysis
|
||||
assert "credentials" in tool._ignore_params
|
||||
|
||||
def test_tool_initialization_without_credentials(self, sample_function):
|
||||
"""Test tool initialization when no credential management is needed.
|
||||
|
||||
Some tools might handle authentication externally or use service
|
||||
accounts, so credential management should be optional.
|
||||
"""
|
||||
tool = BigQueryTool(func=sample_function, credentials_config=None)
|
||||
|
||||
assert tool.func == sample_function
|
||||
assert tool._credentials_manager is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_with_valid_credentials(
|
||||
self, sample_function, credentials_config, mock_tool_context
|
||||
):
|
||||
"""Test successful tool execution with valid credentials.
|
||||
|
||||
This tests the main happy path where credentials are available
|
||||
and the underlying function executes successfully.
|
||||
"""
|
||||
tool = BigQueryTool(
|
||||
func=sample_function, credentials_config=credentials_config
|
||||
)
|
||||
|
||||
# Mock the credentials manager to return valid credentials
|
||||
mock_creds = Mock(spec=Credentials)
|
||||
with patch.object(
|
||||
tool._credentials_manager,
|
||||
"get_valid_credentials",
|
||||
return_value=mock_creds,
|
||||
) as mock_get_creds:
|
||||
|
||||
result = await tool.run_async(
|
||||
args={"param1": "test_value"}, tool_context=mock_tool_context
|
||||
)
|
||||
|
||||
mock_get_creds.assert_called_once_with(mock_tool_context)
|
||||
assert result["result"] == "Success with test_value"
|
||||
assert result["authenticated"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_oauth_flow_in_progress(
|
||||
self, sample_function, credentials_config, mock_tool_context
|
||||
):
|
||||
"""Test tool behavior when OAuth flow is in progress.
|
||||
|
||||
When credentials aren't available and OAuth flow is needed,
|
||||
the tool should return a user-friendly message rather than failing.
|
||||
"""
|
||||
tool = BigQueryTool(
|
||||
func=sample_function, credentials_config=credentials_config
|
||||
)
|
||||
|
||||
# Mock credentials manager to return None (OAuth flow in progress)
|
||||
with patch.object(
|
||||
tool._credentials_manager, "get_valid_credentials", return_value=None
|
||||
) as mock_get_creds:
|
||||
|
||||
result = await tool.run_async(
|
||||
args={"param1": "test_value"}, tool_context=mock_tool_context
|
||||
)
|
||||
|
||||
mock_get_creds.assert_called_once_with(mock_tool_context)
|
||||
assert "authorization is required" in result.lower()
|
||||
assert tool.name in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_without_credentials_manager(
|
||||
self, sample_function, mock_tool_context
|
||||
):
|
||||
"""Test tool execution when no credential management is configured.
|
||||
|
||||
Tools without credential managers should execute normally,
|
||||
passing None for credentials if the function accepts them.
|
||||
"""
|
||||
tool = BigQueryTool(func=sample_function, credentials_config=None)
|
||||
|
||||
result = await tool.run_async(
|
||||
args={"param1": "test_value"}, tool_context=mock_tool_context
|
||||
)
|
||||
|
||||
assert result["result"] == "Success with test_value"
|
||||
assert result["authenticated"] is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_with_async_function(
|
||||
self, async_sample_function, credentials_config, mock_tool_context
|
||||
):
|
||||
"""Test that async functions are properly handled.
|
||||
|
||||
The tool should correctly detect and execute async functions,
|
||||
which is important for tools that make async API calls.
|
||||
"""
|
||||
tool = BigQueryTool(
|
||||
func=async_sample_function, credentials_config=credentials_config
|
||||
)
|
||||
|
||||
mock_creds = Mock(spec=Credentials)
|
||||
with patch.object(
|
||||
tool._credentials_manager,
|
||||
"get_valid_credentials",
|
||||
return_value=mock_creds,
|
||||
):
|
||||
|
||||
result = await tool.run_async(
|
||||
args={"param1": "test_value"}, tool_context=mock_tool_context
|
||||
)
|
||||
|
||||
assert result["result"] == "Async success with test_value"
|
||||
assert result["authenticated"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_exception_handling(
|
||||
self, credentials_config, mock_tool_context
|
||||
):
|
||||
"""Test that exceptions in tool execution are properly handled.
|
||||
|
||||
Tools should gracefully handle errors and return structured
|
||||
error responses rather than letting exceptions propagate.
|
||||
"""
|
||||
|
||||
def failing_function(param1: str, credentials: Credentials = None) -> dict:
|
||||
raise ValueError("Something went wrong")
|
||||
|
||||
tool = BigQueryTool(
|
||||
func=failing_function, credentials_config=credentials_config
|
||||
)
|
||||
|
||||
mock_creds = Mock(spec=Credentials)
|
||||
with patch.object(
|
||||
tool._credentials_manager,
|
||||
"get_valid_credentials",
|
||||
return_value=mock_creds,
|
||||
):
|
||||
|
||||
result = await tool.run_async(
|
||||
args={"param1": "test_value"}, tool_context=mock_tool_context
|
||||
)
|
||||
|
||||
assert result["status"] == "ERROR"
|
||||
assert "Something went wrong" in result["error_details"]
|
||||
|
||||
def test_function_signature_analysis(self, credentials_config):
|
||||
"""Test that function signature analysis correctly handles credentials parameter.
|
||||
|
||||
The tool should properly identify and handle the credentials parameter
|
||||
while preserving other parameter analysis for LLM function calling.
|
||||
"""
|
||||
|
||||
def complex_function(
|
||||
required_param: str,
|
||||
optional_param: str = "default",
|
||||
credentials: Credentials = None,
|
||||
) -> dict:
|
||||
return {"success": True}
|
||||
|
||||
tool = BigQueryTool(
|
||||
func=complex_function, credentials_config=credentials_config
|
||||
)
|
||||
|
||||
# The 'credentials' parameter should be ignored in mandatory args analysis
|
||||
mandatory_args = tool._get_mandatory_args()
|
||||
assert "required_param" in mandatory_args
|
||||
assert "credentials" not in mandatory_args
|
||||
assert "optional_param" not in mandatory_args
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"input_config, expected_config",
|
||||
[
|
||||
pytest.param(
|
||||
BigQueryToolConfig(
|
||||
write_mode="blocked", max_query_result_rows=50
|
||||
),
|
||||
BigQueryToolConfig(
|
||||
write_mode="blocked", max_query_result_rows=50
|
||||
),
|
||||
id="with_provided_config",
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
BigQueryToolConfig(),
|
||||
id="with_none_config_creates_default",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_tool_config_initialization(self, input_config, expected_config):
|
||||
"""Tests that self._tool_config is correctly initialized by comparing its
|
||||
|
||||
final state to an expected configuration object.
|
||||
"""
|
||||
# 1. Initialize the tool with the parameterized config
|
||||
tool = BigQueryTool(func=None, bigquery_tool_config=input_config)
|
||||
|
||||
# 2. Assert that the tool's config has the same attribute values
|
||||
# as the expected config. Comparing the __dict__ is a robust
|
||||
# way to check for value equality.
|
||||
assert tool._tool_config.__dict__ == expected_config.__dict__ # pylint: disable=protected-access
|
||||
@@ -15,8 +15,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from google.adk.tools.bigquery import BigQueryCredentialsConfig
|
||||
from google.adk.tools.bigquery import BigQueryTool
|
||||
from google.adk.tools.bigquery import BigQueryToolset
|
||||
from google.adk.tools.bigquery.config import BigQueryToolConfig
|
||||
from google.adk.tools.google_tool import GoogleTool
|
||||
import pytest
|
||||
|
||||
|
||||
@@ -30,12 +31,18 @@ async def test_bigquery_toolset_tools_default():
|
||||
credentials_config = BigQueryCredentialsConfig(
|
||||
client_id="abc", client_secret="def"
|
||||
)
|
||||
toolset = BigQueryToolset(credentials_config=credentials_config)
|
||||
toolset = BigQueryToolset(
|
||||
credentials_config=credentials_config, bigquery_tool_config=None
|
||||
)
|
||||
# Verify that the tool config is initialized to default values.
|
||||
assert isinstance(toolset._tool_settings, BigQueryToolConfig) # pylint: disable=protected-access
|
||||
assert toolset._tool_settings.__dict__ == BigQueryToolConfig().__dict__ # pylint: disable=protected-access
|
||||
|
||||
tools = await toolset.get_tools()
|
||||
assert tools is not None
|
||||
|
||||
assert len(tools) == 5
|
||||
assert all([isinstance(tool, BigQueryTool) for tool in tools])
|
||||
assert all([isinstance(tool, GoogleTool) for tool in tools])
|
||||
|
||||
expected_tool_names = set([
|
||||
"list_dataset_ids",
|
||||
@@ -77,7 +84,7 @@ async def test_bigquery_toolset_tools_selective(selected_tools):
|
||||
assert tools is not None
|
||||
|
||||
assert len(tools) == len(selected_tools)
|
||||
assert all([isinstance(tool, BigQueryTool) for tool in tools])
|
||||
assert all([isinstance(tool, GoogleTool) for tool in tools])
|
||||
|
||||
expected_tool_names = set(selected_tools)
|
||||
actual_tool_names = set([tool.name for tool in tools])
|
||||
@@ -114,7 +121,7 @@ async def test_bigquery_toolset_unknown_tool(selected_tools, returned_tools):
|
||||
assert tools is not None
|
||||
|
||||
assert len(tools) == len(returned_tools)
|
||||
assert all([isinstance(tool, BigQueryTool) for tool in tools])
|
||||
assert all([isinstance(tool, GoogleTool) for tool in tools])
|
||||
|
||||
expected_tool_names = set(returned_tools)
|
||||
actual_tool_names = set([tool.name for tool in tools])
|
||||
|
||||
Reference in New Issue
Block a user