mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: set per-tool user agent in BQ calls and tool label in BQ jobs
This will help per tool usage for BigQuery tools. PiperOrigin-RevId: 829142106
This commit is contained in:
committed by
Copybara-Service
parent
f1f44675e4
commit
c0be1df052
@@ -156,6 +156,31 @@ def test_bigquery_client_user_agent_custom():
|
||||
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(
|
||||
"google.cloud.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
|
||||
|
||||
Reference in New Issue
Block a user