@@ -1673,7 +1673,9 @@ def test_execute_sql_job_labels(
16731673 query = "SELECT 123 AS num"
16741674 statement_type = "SELECT"
16751675 credentials = mock .create_autospec (Credentials , instance = True )
1676- tool_settings = BigQueryToolConfig (write_mode = write_mode )
1676+ tool_settings = BigQueryToolConfig (
1677+ write_mode = write_mode , application_name = "test-app"
1678+ )
16771679 tool_context = mock .create_autospec (ToolContext , instance = True )
16781680 tool_context .state .get .return_value = None
16791681
@@ -1702,12 +1704,13 @@ def test_execute_sql_job_labels(
17021704 for call_args in call_args_list :
17031705 _ , mock_kwargs = call_args
17041706 assert mock_kwargs ["job_config" ].labels == {
1705- "adk-bigquery-tool" : "execute_sql"
1707+ "adk-bigquery-tool" : "execute_sql" ,
1708+ "adk-bigquery-application-name" : "test-app" ,
17061709 }
17071710
17081711
17091712@pytest .mark .parametrize (
1710- ("tool_call" , "expected_label " ),
1713+ ("tool_call" , "expected_tool_label " ),
17111714 [
17121715 pytest .param (
17131716 lambda tool_context : query_tool .forecast (
@@ -1751,7 +1754,7 @@ def test_execute_sql_job_labels(
17511754 ),
17521755 ],
17531756)
1754- def test_ml_tool_job_labels (tool_call , expected_label ):
1757+ def test_ml_tool_job_labels (tool_call , expected_tool_label ):
17551758 """Test ML tools for job label."""
17561759
17571760 with mock .patch .object (bigquery , "Client" , autospec = True ) as Client :
@@ -1768,10 +1771,85 @@ def test_ml_tool_job_labels(tool_call, expected_label):
17681771 for call_args in call_args_list :
17691772 _ , mock_kwargs = call_args
17701773 assert mock_kwargs ["job_config" ].labels == {
1771- "adk-bigquery-tool" : expected_label
1774+ "adk-bigquery-tool" : expected_tool_label
17721775 }
17731776
17741777
1778+ @pytest .mark .parametrize (
1779+ ("tool_call" , "expected_tool_label" ),
1780+ [
1781+ pytest .param (
1782+ lambda tool_context : query_tool .forecast (
1783+ project_id = "test-project" ,
1784+ history_data = "SELECT * FROM `test-dataset.test-table`" ,
1785+ timestamp_col = "ts_col" ,
1786+ data_col = "data_col" ,
1787+ credentials = mock .create_autospec (Credentials , instance = True ),
1788+ settings = BigQueryToolConfig (
1789+ write_mode = WriteMode .ALLOWED , application_name = "test-app"
1790+ ),
1791+ tool_context = tool_context ,
1792+ ),
1793+ "forecast" ,
1794+ id = "forecast-app-name" ,
1795+ ),
1796+ pytest .param (
1797+ lambda tool_context : query_tool .analyze_contribution (
1798+ project_id = "test-project" ,
1799+ input_data = "test-dataset.test-table" ,
1800+ dimension_id_cols = ["dim1" , "dim2" ],
1801+ contribution_metric = "SUM(metric)" ,
1802+ is_test_col = "is_test" ,
1803+ credentials = mock .create_autospec (Credentials , instance = True ),
1804+ settings = BigQueryToolConfig (
1805+ write_mode = WriteMode .ALLOWED , application_name = "test-app"
1806+ ),
1807+ tool_context = tool_context ,
1808+ ),
1809+ "analyze_contribution" ,
1810+ id = "analyze-contribution-app-name" ,
1811+ ),
1812+ pytest .param (
1813+ lambda tool_context : query_tool .detect_anomalies (
1814+ project_id = "test-project" ,
1815+ history_data = "SELECT * FROM `test-dataset.test-table`" ,
1816+ times_series_timestamp_col = "ts_timestamp" ,
1817+ times_series_data_col = "ts_data" ,
1818+ credentials = mock .create_autospec (Credentials , instance = True ),
1819+ settings = BigQueryToolConfig (
1820+ write_mode = WriteMode .ALLOWED , application_name = "test-app"
1821+ ),
1822+ tool_context = tool_context ,
1823+ ),
1824+ "detect_anomalies" ,
1825+ id = "detect-anomalies-app-name" ,
1826+ ),
1827+ ],
1828+ )
1829+ def test_ml_tool_job_labels_w_application_name (tool_call , expected_tool_label ):
1830+ """Test ML tools for job label with application name."""
1831+
1832+ with mock .patch .object (bigquery , "Client" , autospec = True ) as Client :
1833+ bq_client = Client .return_value
1834+
1835+ tool_context = mock .create_autospec (ToolContext , instance = True )
1836+ tool_context .state .get .return_value = None
1837+ tool_call (tool_context )
1838+
1839+ expected_labels = {
1840+ "adk-bigquery-tool" : expected_tool_label ,
1841+ "adk-bigquery-application-name" : "test-app" ,
1842+ }
1843+
1844+ for call_args_list in [
1845+ bq_client .query .call_args_list ,
1846+ bq_client .query_and_wait .call_args_list ,
1847+ ]:
1848+ for call_args in call_args_list :
1849+ _ , mock_kwargs = call_args
1850+ assert mock_kwargs ["job_config" ].labels == expected_labels
1851+
1852+
17751853def test_execute_sql_max_rows_config ():
17761854 """Test execute_sql tool respects max_query_result_rows from config."""
17771855 project = "my_project"
0 commit comments