mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
2b0f953255
commit
fd33610e96
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user