test: add tests for max_query_result_rows in BigQuery tool config

Merge https://github.com/google/adk-python/pull/2167, during which:

- Let the already added `max_query_result_rows` field cover for `max_downloaded_rows` field added in the PR
- Revert `max_rows` parameter added to execute_sql function, the tool config control should serve most practical use cases
- Keep the relevant tests for tool config

🤖 Generated with [Claude Code](https://claude.ai/code)

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/2167 from lupuletic:feature/configurable-max-downloaded-rows 23c56905c297d7aec2be4f1eb86ea23c8178bf21
PiperOrigin-RevId: 830701093
This commit is contained in:
Catalin Lupuleti
2025-11-10 19:43:24 -08:00
committed by Copybara-Service
parent 2b0f953255
commit fd33610e96
2 changed files with 67 additions and 0 deletions
@@ -1773,3 +1773,56 @@ def test_ml_tool_job_labels(tool_call, expected_label):
assert mock_kwargs["job_config"].labels == {
"adk-bigquery-tool": expected_label
}
def test_execute_sql_max_rows_config():
"""Test execute_sql tool respects max_query_result_rows from config."""
project = "my_project"
query = "SELECT 123 AS num"
statement_type = "SELECT"
query_result = [{"num": i} for i in range(20)] # 20 rows
credentials = mock.create_autospec(Credentials, instance=True)
tool_config = BigQueryToolConfig(max_query_result_rows=10)
tool_context = mock.create_autospec(ToolContext, instance=True)
with mock.patch("google.cloud.bigquery.Client", autospec=False) as Client:
bq_client = Client.return_value
query_job = mock.create_autospec(bigquery.QueryJob)
query_job.statement_type = statement_type
bq_client.query.return_value = query_job
bq_client.query_and_wait.return_value = query_result[:10]
result = execute_sql(project, query, credentials, tool_config, tool_context)
# Check that max_results was called with config value
bq_client.query_and_wait.assert_called_once()
call_args = bq_client.query_and_wait.call_args
assert call_args.kwargs["max_results"] == 10
# Check truncation flag is set
assert result["status"] == "SUCCESS"
assert result["result_is_likely_truncated"] is True
def test_execute_sql_no_truncation():
"""Test execute_sql tool when results are not truncated."""
project = "my_project"
query = "SELECT 123 AS num"
statement_type = "SELECT"
query_result = [{"num": i} for i in range(3)] # Only 3 rows
credentials = mock.create_autospec(Credentials, instance=True)
tool_config = BigQueryToolConfig(max_query_result_rows=10)
tool_context = mock.create_autospec(ToolContext, instance=True)
with mock.patch("google.cloud.bigquery.Client", autospec=False) as Client:
bq_client = Client.return_value
query_job = mock.create_autospec(bigquery.QueryJob)
query_job.statement_type = statement_type
bq_client.query.return_value = query_job
bq_client.query_and_wait.return_value = query_result
result = execute_sql(project, query, credentials, tool_config, tool_context)
# Check no truncation flag when fewer rows than limit
assert result["status"] == "SUCCESS"
assert "result_is_likely_truncated" not in result
@@ -42,3 +42,17 @@ def test_bigquery_tool_config_invalid_application_name():
match="Application name should not contain spaces.",
):
BigQueryToolConfig(application_name="my agent")
def test_bigquery_tool_config_max_query_result_rows_default():
"""Test BigQueryToolConfig max_query_result_rows default value."""
with pytest.warns(UserWarning):
config = BigQueryToolConfig()
assert config.max_query_result_rows == 50
def test_bigquery_tool_config_max_query_result_rows_custom():
"""Test BigQueryToolConfig max_query_result_rows custom value."""
with pytest.warns(UserWarning):
config = BigQueryToolConfig(max_query_result_rows=100)
assert config.max_query_result_rows == 100