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