From fd33610e967ad814bc02422f5d14dae046bee833 Mon Sep 17 00:00:00 2001
From: Catalin Lupuleti <105351510+lupuletic@users.noreply.github.com>
Date: Mon, 10 Nov 2025 19:42:55 -0800
Subject: [PATCH] test: add tests for `max_query_result_rows` in BigQuery tool
 config
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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
---
 .../bigquery/test_bigquery_query_tool.py      | 53 +++++++++++++++++++
 .../bigquery/test_bigquery_tool_config.py     | 14 +++++
 2 files changed, 67 insertions(+)

diff --git a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py b/tests/unittests/tools/bigquery/test_bigquery_query_tool.py
index d4906e3d..547099b6 100644
--- a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py
+++ b/tests/unittests/tools/bigquery/test_bigquery_query_tool.py
@@ -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
diff --git a/tests/unittests/tools/bigquery/test_bigquery_tool_config.py b/tests/unittests/tools/bigquery/test_bigquery_tool_config.py
index 1552064c..19b81c00 100644
--- a/tests/unittests/tools/bigquery/test_bigquery_tool_config.py
+++ b/tests/unittests/tools/bigquery/test_bigquery_tool_config.py
@@ -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